Skip to content

Commit d458d60

Browse files
chernserchernser
authored
Merge pull request #1904 from ClickHouse/v2_jwt_auth
[client-v2] Added implementation for Bearer token auth
2 parents 19172f6 + da03749 commit d458d60

File tree

5 files changed

+139
-13
lines changed

5 files changed

+139
-13
lines changed

‎client-v2/src/main/java/com/clickhouse/client/api/Client.java‎

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,6 @@
118118
* ...
119119
* }
120120
* }
121-
*
122121
* }
123122
*
124123
*
@@ -132,6 +131,9 @@ public class Client implements AutoCloseable {
132131

133132
private final Set<String> endpoints;
134133
private final Map<String, String> configuration;
134+
135+
private final Map<String, String> readOnlyConfig;
136+
135137
private final List<ClickHouseNode> serverNodes = new ArrayList<>();
136138

137139
// POJO serializer mapping (class -> (schema -> (format -> serializer)))
@@ -158,6 +160,7 @@ private Client(Set<String> endpoints, Map<String,String> configuration, boolean
158160
ExecutorService sharedOperationExecutor, ColumnToMethodMatchingStrategy columnToMethodMatchingStrategy) {
159161
this.endpoints = endpoints;
160162
this.configuration = configuration;
163+
this.readOnlyConfig = Collections.unmodifiableMap(this.configuration);
161164
this.endpoints.forEach(endpoint -> {
162165
this.serverNodes.add(ClickHouseNode.of(endpoint, this.configuration));
163166
});
@@ -853,7 +856,7 @@ public Builder allowBinaryReaderToReuseBuffers(boolean reuse) {
853856
* @return same instance of the builder
854857
*/
855858
public Builder httpHeader(String key, String value) {
856-
this.configuration.put(ClientConfigProperties.HTTP_HEADER_PREFIX + key.toUpperCase(Locale.US), value);
859+
this.configuration.put(ClientConfigProperties.httpHeader(key), value);
857860
return this;
858861
}
859862

@@ -864,7 +867,7 @@ public Builder httpHeader(String key, String value) {
864867
* @return same instance of the builder
865868
*/
866869
public Builder httpHeader(String key, Collection<String> values) {
867-
this.configuration.put(ClientConfigProperties.HTTP_HEADER_PREFIX + key.toUpperCase(Locale.US), ClientConfigProperties.commaSeparated(values));
870+
this.configuration.put(ClientConfigProperties.httpHeader(key), ClientConfigProperties.commaSeparated(values));
868871
return this;
869872
}
870873

@@ -955,6 +958,19 @@ public Builder setOptions(Map<String, String> options) {
955958
return this;
956959
}
957960

961+
/**
962+
* Specifies whether to use Bearer Authentication and what token to use.
963+
* The token will be sent as is, so it should be encoded before passing to this method.
964+
*
965+
* @param bearerToken - token to use
966+
* @return same instance of the builder
967+
*/
968+
public Builder useBearerTokenAuth(String bearerToken) {
969+
// Most JWT libraries (https://jwt.io/libraries?language=Java) compact tokens in proper way
970+
this.httpHeader(HttpHeaders.AUTHORIZATION, "Bearer " + bearerToken);
971+
return this;
972+
}
973+
958974
public Client build() {
959975
setDefaults();
960976

@@ -965,8 +981,9 @@ public Client build() {
965981
// check if username and password are empty. so can not initiate client?
966982
if (!this.configuration.containsKey("access_token") &&
967983
(!this.configuration.containsKey("user") || !this.configuration.containsKey("password")) &&
968-
!MapUtils.getFlag(this.configuration, "ssl_authentication", false)) {
969-
throw new IllegalArgumentException("Username and password (or access token, or SSL authentication) are required");
984+
!MapUtils.getFlag(this.configuration, "ssl_authentication", false) &&
985+
!this.configuration.containsKey(ClientConfigProperties.httpHeader(HttpHeaders.AUTHORIZATION))) {
986+
throw new IllegalArgumentException("Username and password (or access token or SSL authentication or pre-define Authorization header) are required");
970987
}
971988

972989
if (this.configuration.containsKey("ssl_authentication") &&
@@ -1012,7 +1029,8 @@ public Client build() {
10121029
throw new IllegalArgumentException("Nor server timezone nor specific timezone is set");
10131030
}
10141031

1015-
return new Client(this.endpoints, this.configuration, this.useNewImplementation, this.sharedOperationExecutor, this.columnToMethodMatchingStrategy);
1032+
return new Client(this.endpoints, this.configuration, this.useNewImplementation, this.sharedOperationExecutor,
1033+
this.columnToMethodMatchingStrategy);
10161034
}
10171035

10181036
private static final int DEFAULT_NETWORK_BUFFER_SIZE = 300_000;
@@ -2104,7 +2122,7 @@ public String toString() {
21042122
* @return - configuration options
21052123
*/
21062124
public Map<String, String> getConfiguration() {
2107-
return Collections.unmodifiableMap(configuration);
2125+
return readOnlyConfig;
21082126
}
21092127

21102128
/** Returns operation timeout in seconds */
@@ -2151,6 +2169,10 @@ public Collection<String> getDBRoles() {
21512169
return unmodifiableDbRolesView;
21522170
}
21532171

2172+
public void updateBearerToken(String bearer) {
2173+
this.configuration.put(ClientConfigProperties.httpHeader(HttpHeaders.AUTHORIZATION), "Bearer " + bearer);
2174+
}
2175+
21542176
private ClickHouseNode getNextAliveNode() {
21552177
return serverNodes.get(0);
21562178
}

‎client-v2/src/main/java/com/clickhouse/client/api/ClientConfigProperties.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import java.util.Collection;
55
import java.util.Collections;
66
import java.util.List;
7+
import java.util.Locale;
78
import java.util.stream.Collectors;
89

910
/**
@@ -157,6 +158,10 @@ public static String serverSetting(String key) {
157158
return SERVER_SETTING_PREFIX + key;
158159
}
159160

161+
public static String httpHeader(String key) {
162+
return HTTP_HEADER_PREFIX + key.toUpperCase(Locale.US);
163+
}
164+
160165
public static String commaSeparated(Collection<?> values) {
161166
StringBuilder sb = new StringBuilder();
162167
for (Object value : values) {

‎client-v2/src/main/java/com/clickhouse/client/api/ServerException.java‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,34 @@ public class ServerException extends RuntimeException {
77
public static final int TABLE_NOT_FOUND = 60;
88

99
private final int code;
10+
11+
private final int transportProtocolCode;
12+
1013
public ServerException(int code, String message) {
14+
this(code, message, 500);
15+
}
16+
17+
public ServerException(int code, String message, int transportProtocolCode) {
1118
super(message);
1219
this.code = code;
20+
this.transportProtocolCode = transportProtocolCode;
1321
}
1422

23+
/**
24+
* Returns CH server error code. May return 0 if code is unknown.
25+
* @return - error code from server response
26+
*/
1527
public int getCode() {
1628
return code;
1729
}
30+
31+
/**
32+
* Returns error code of underlying transport protocol. For example, HTTP status.
33+
* By default, will return {@code 500 } what is derived from HTTP Server Internal Error.
34+
*
35+
* @return - transport status code
36+
*/
37+
public int getTransportProtocolCode() {
38+
return transportProtocolCode;
39+
}
1840
}

‎client-v2/src/main/java/com/clickhouse/client/api/internal/HttpAPIClientHelper.java‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@
7373
import java.util.Set;
7474
import java.util.concurrent.TimeUnit;
7575
import java.util.function.Function;
76+
import java.util.function.Supplier;
7677

7778
import static com.clickhouse.client.api.ClientConfigProperties.SOCKET_TCP_NO_DELAY_OPT;
7879

@@ -335,10 +336,13 @@ public Exception readError(ClassicHttpResponse httpResponse) {
335336

336337
String msg = msgBuilder.toString().replaceAll("\\s+", " ").replaceAll("\\\\n", " ")
337338
.replaceAll("\\\\/", "/");
338-
return new ServerException(serverCode, msg);
339+
if (msg.trim().isEmpty()) {
340+
msg = String.format(ERROR_CODE_PREFIX_PATTERN, serverCode) + " <Unreadable error message> (transport error: " + httpResponse.getCode() + ")";
341+
}
342+
return new ServerException(serverCode, msg, httpResponse.getCode());
339343
} catch (Exception e) {
340344
LOG.error("Failed to read error message", e);
341-
return new ServerException(serverCode, String.format(ERROR_CODE_PREFIX_PATTERN, serverCode) + " <Unreadable error message>");
345+
return new ServerException(serverCode, String.format(ERROR_CODE_PREFIX_PATTERN, serverCode) + " <Unreadable error message> (transport error: " + httpResponse.getCode() + ")", httpResponse.getCode());
342346
}
343347
}
344348

@@ -450,12 +454,12 @@ private void addHeaders(HttpPost req, Map<String, String> chConfig, Map<String,
450454

451455
for (Map.Entry<String, String> entry : chConfig.entrySet()) {
452456
if (entry.getKey().startsWith(ClientConfigProperties.HTTP_HEADER_PREFIX)) {
453-
req.addHeader(entry.getKey().substring(ClientConfigProperties.HTTP_HEADER_PREFIX.length()), entry.getValue());
457+
req.setHeader(entry.getKey().substring(ClientConfigProperties.HTTP_HEADER_PREFIX.length()), entry.getValue());
454458
}
455459
}
456460
for (Map.Entry<String, Object> entry : requestConfig.entrySet()) {
457461
if (entry.getKey().startsWith(ClientConfigProperties.HTTP_HEADER_PREFIX)) {
458-
req.addHeader(entry.getKey().substring(ClientConfigProperties.HTTP_HEADER_PREFIX.length()), entry.getValue().toString());
462+
req.setHeader(entry.getKey().substring(ClientConfigProperties.HTTP_HEADER_PREFIX.length()), entry.getValue().toString());
459463
}
460464
}
461465

‎client-v2/src/test/java/com/clickhouse/client/HttpTransportTests.java‎

Lines changed: 75 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
import java.io.ByteArrayInputStream;
3939
import java.net.Socket;
4040
import java.nio.ByteBuffer;
41+
import java.nio.charset.StandardCharsets;
4142
import java.time.temporal.ChronoUnit;
4243
import java.util.Arrays;
4344
import java.util.Base64;
@@ -51,8 +52,11 @@
5152
import java.util.concurrent.TimeUnit;
5253
import java.util.concurrent.atomic.AtomicInteger;
5354
import java.util.regex.Pattern;
55+
import java.util.function.Supplier;
5456

5557
import static com.github.tomakehurst.wiremock.stubbing.Scenario.STARTED;
58+
import static org.junit.Assert.assertEquals;
59+
import static org.junit.Assert.assertThrows;
5660
import static org.junit.Assert.fail;
5761

5862
public class HttpTransportTests extends BaseIntegrationTest {
@@ -66,7 +70,6 @@ public void testConnectionTTL(Long connectionTtl, Long keepAlive, int openSocket
6670
ClickHouseNode server = getServer(ClickHouseProtocol.HTTP);
6771

6872
int proxyPort = new Random().nextInt(1000) + 10000;
69-
System.out.println("proxyPort: " + proxyPort);
7073
ConnectionCounterListener connectionCounter = new ConnectionCounterListener();
7174
WireMockServer proxy = new WireMockServer(WireMockConfiguration
7275
.options().port(proxyPort)
@@ -154,7 +157,6 @@ public void closed(Socket socket) {
154157
public void testConnectionRequestTimeout() {
155158

156159
int serverPort = new Random().nextInt(1000) + 10000;
157-
System.out.println("proxyPort: " + serverPort);
158160
ConnectionCounterListener connectionCounter = new ConnectionCounterListener();
159161
WireMockServer proxy = new WireMockServer(WireMockConfiguration
160162
.options().port(serverPort)
@@ -794,4 +796,75 @@ public static Object[][] testUserAgentHasCompleteProductName_dataProvider() {
794796
{ "test-client/1.0", Pattern.compile("test-client/1.0 clickhouse-java-v2\\/.+ \\(.+\\) Apache HttpClient\\/[\\d\\.]+$")},
795797
{ "test-client/", Pattern.compile("test-client/ clickhouse-java-v2\\/.+ \\(.+\\) Apache HttpClient\\/[\\d\\.]+$")}};
796798
}
799+
800+
@Test(groups = { "integration" })
801+
public void testBearerTokenAuth() throws Exception {
802+
WireMockServer mockServer = new WireMockServer( WireMockConfiguration
803+
.options().port(9090).notifier(new ConsoleNotifier(false)));
804+
mockServer.start();
805+
806+
try {
807+
String jwtToken1 = Arrays.stream(
808+
new String[]{"header", "payload", "signature"})
809+
.map(s -> Base64.getEncoder().encodeToString(s.getBytes(StandardCharsets.UTF_8)))
810+
.reduce((s1, s2) -> s1 + "." + s2).get();
811+
try (Client client = new Client.Builder().addEndpoint(Protocol.HTTP, "localhost", mockServer.port(), false)
812+
.useBearerTokenAuth(jwtToken1)
813+
.compressServerResponse(false)
814+
.build()) {
815+
816+
mockServer.addStubMapping(WireMock.post(WireMock.anyUrl())
817+
.withHeader("Authorization", WireMock.equalTo("Bearer " + jwtToken1))
818+
.willReturn(WireMock.aResponse()
819+
.withHeader("X-ClickHouse-Summary",
820+
"{ \"read_bytes\": \"10\", \"read_rows\": \"1\"}")).build());
821+
822+
try (QueryResponse response = client.query("SELECT 1").get(1, TimeUnit.SECONDS)) {
823+
Assert.assertEquals(response.getReadBytes(), 10);
824+
} catch (Exception e) {
825+
Assert.fail("Unexpected exception", e);
826+
}
827+
}
828+
829+
String jwtToken2 = Arrays.stream(
830+
new String[]{"header2", "payload2", "signature2"})
831+
.map(s -> Base64.getEncoder().encodeToString(s.getBytes(StandardCharsets.UTF_8)))
832+
.reduce((s1, s2) -> s1 + "." + s2).get();
833+
834+
mockServer.resetAll();
835+
mockServer.addStubMapping(WireMock.post(WireMock.anyUrl())
836+
.withHeader("Authorization", WireMock.equalTo("Bearer " + jwtToken1))
837+
.willReturn(WireMock.aResponse()
838+
.withStatus(HttpStatus.SC_UNAUTHORIZED))
839+
.build());
840+
841+
try (Client client = new Client.Builder().addEndpoint(Protocol.HTTP, "localhost", mockServer.port(), false)
842+
.useBearerTokenAuth(jwtToken1)
843+
.compressServerResponse(false)
844+
.build()) {
845+
846+
try {
847+
client.execute("SELECT 1").get();
848+
fail("Exception expected");
849+
} catch (ServerException e) {
850+
Assert.assertEquals(e.getTransportProtocolCode(), HttpStatus.SC_UNAUTHORIZED);
851+
}
852+
853+
mockServer.resetAll();
854+
mockServer.addStubMapping(WireMock.post(WireMock.anyUrl())
855+
.withHeader("Authorization", WireMock.equalTo("Bearer " + jwtToken2))
856+
.willReturn(WireMock.aResponse()
857+
.withHeader("X-ClickHouse-Summary",
858+
"{ \"read_bytes\": \"10\", \"read_rows\": \"1\"}"))
859+
860+
.build());
861+
862+
client.updateBearerToken(jwtToken2);
863+
864+
client.execute("SELECT 1").get();
865+
}
866+
} finally {
867+
mockServer.stop();
868+
}
869+
}
797870
}

0 commit comments

Comments
 (0)