fix: add Edge authentication support (#158)

alespour · web-flow · commit b67f42364175 · 2024-07-22T17:01:55.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,7 @@
 
 1. [#148](https://github.com/InfluxCommunity/influxdb3-java/pull/148): InfluxDB Edge (OSS) error handling
 1. [#153](https://github.com/InfluxCommunity/influxdb3-java/pull/153): Parsing timestamp columns
+1. [#158](https://github.com/InfluxCommunity/influxdb3-java/pull/158): Add InfluxDB Edge (OSS) authentication support.
 
 ## 0.8.0 [2024-06-24]
 
diff --git a/src/main/java/com/influxdb/v3/client/InfluxDBClient.java b/src/main/java/com/influxdb/v3/client/InfluxDBClient.java
@@ -444,6 +444,7 @@ static InfluxDBClient getInstance(@Nonnull final String connectionString) {
      * <ul>
      *   <li>INFLUX_HOST - cloud/server URL <i>required</i></li>
      *   <li>INFLUX_TOKEN - authentication token <i>required</i></li>
+     *   <li>INFLUX_AUTH_SCHEME - authentication scheme</li>
      *   <li>INFLUX_ORG - organization name</li>
      *   <li>INFLUX_DATABASE - database (bucket) name</li>
      *   <li>INFLUX_PRECISION - timestamp precision when writing data</li>
@@ -453,6 +454,7 @@ static InfluxDBClient getInstance(@Nonnull final String connectionString) {
      * <ul>
      *   <li>influx.host - cloud/server URL <i>required</i></li>
      *   <li>influx.token - authentication token <i>required</i></li>
+     *   <li>influx.authScheme - authentication scheme</li>
      *   <li>influx.org - organization name</li>
      *   <li>influx.database - database (bucket) name</li>
      *   <li>influx.precision - timestamp precision when writing data</li>
diff --git a/src/main/java/com/influxdb/v3/client/config/ClientConfig.java b/src/main/java/com/influxdb/v3/client/config/ClientConfig.java
@@ -46,6 +46,7 @@
  * <ul>
  *     <li><code>host</code> - hostname or IP address of the InfluxDB server</li>
  *     <li><code>token</code> - authentication token for accessing the InfluxDB server</li>
+ *     <li><code>authScheme</code> - authentication scheme</li>
  *     <li><code>organization</code> - organization to be used for operations</li>
  *     <li><code>database</code> - database to be used for InfluxDB operations</li>
  *     <li><code>writePrecision</code> - precision to use when writing points to InfluxDB</li>
@@ -86,6 +87,7 @@ public final class ClientConfig {
 
     private final String host;
     private final char[] token;
+    private final String authScheme;
     private final String organization;
     private final String database;
     private final WritePrecision writePrecision;
@@ -118,6 +120,16 @@ public char[] getToken() {
         return token;
     }
 
+    /**
+     * Gets authentication scheme.
+     *
+     * @return authentication scheme, may be null
+     */
+    @Nullable
+    public String getAuthScheme() {
+        return authScheme;
+    }
+
     /**
      * Gets organization to be used for operations.
      *
@@ -247,6 +259,7 @@ public boolean equals(final Object o) {
         ClientConfig that = (ClientConfig) o;
         return Objects.equals(host, that.host)
                 && Arrays.equals(token, that.token)
+                && Objects.equals(authScheme, that.authScheme)
                 && Objects.equals(organization, that.organization)
                 && Objects.equals(database, that.database)
                 && writePrecision == that.writePrecision
@@ -262,7 +275,7 @@ public boolean equals(final Object o) {
 
     @Override
     public int hashCode() {
-        return Objects.hash(host, Arrays.hashCode(token), organization,
+        return Objects.hash(host, Arrays.hashCode(token), authScheme, organization,
           database, writePrecision, gzipThreshold,
           timeout, allowHttpRedirects, disableServerCertificateValidation,
           proxy, authenticator, headers,
@@ -295,6 +308,7 @@ public String toString() {
     public static final class Builder {
         private String host;
         private char[] token;
+        private String authScheme;
         private String organization;
         private String database;
         private WritePrecision writePrecision;
@@ -333,6 +347,20 @@ public Builder token(@Nullable final char[] token) {
             return this;
         }
 
+        /**
+         * Sets authentication scheme.
+         *
+         * @param authScheme authentication scheme.
+         *                   Default <code>null</code> for Cloud access, set to 'Bearer' for Edge.
+         * @return this
+         */
+        @Nonnull
+        public Builder authScheme(@Nullable final String authScheme) {
+
+            this.authScheme = authScheme;
+            return this;
+        }
+
         /**
          * Sets organization to be used for operations.
          *
@@ -525,6 +553,9 @@ public ClientConfig build(@Nonnull final String connectionString) throws Malform
             if (parameters.containsKey("token")) {
                 this.token(parameters.get("token").toCharArray());
             }
+            if (parameters.containsKey("authScheme")) {
+                this.authScheme(parameters.get("authScheme"));
+            }
             if (parameters.containsKey("org")) {
                 this.organization(parameters.get("org"));
             }
@@ -565,6 +596,10 @@ public ClientConfig build(@Nonnull final Map<String, String> env, final Properti
             if (token != null) {
                 this.token(token.toCharArray());
             }
+            final String authScheme = get.apply("INFLUX_AUTH_SCHEME", "influx.authScheme");
+            if (authScheme != null) {
+                this.authScheme(authScheme);
+            }
             final String org = get.apply("INFLUX_ORG", "influx.org");
             if (org != null) {
                 this.organization(org);
@@ -612,6 +647,7 @@ private WritePrecision parsePrecision(@Nonnull final String precision) {
     private ClientConfig(@Nonnull final Builder builder) {
         host = builder.host;
         token = builder.token;
+        authScheme = builder.authScheme;
         organization = builder.organization;
         database = builder.database;
         writePrecision = builder.writePrecision != null ? builder.writePrecision : WritePrecision.NS;
diff --git a/src/main/java/com/influxdb/v3/client/internal/RestClient.java b/src/main/java/com/influxdb/v3/client/internal/RestClient.java
@@ -162,7 +162,11 @@ void request(@Nonnull final String path,
         }
         request.header("User-Agent", userAgent);
         if (config.getToken() != null && config.getToken().length > 0) {
-            request.header("Authorization", String.format("Token %s", new String(config.getToken())));
+            String authScheme = config.getAuthScheme();
+            if (authScheme == null) {
+                authScheme = "Token";
+            }
+            request.header("Authorization", String.format("%s %s", authScheme, new String(config.getToken())));
         }
 
         HttpResponse<String> response;
diff --git a/src/test/java/com/influxdb/v3/client/config/ClientConfigTest.java b/src/test/java/com/influxdb/v3/client/config/ClientConfigTest.java
@@ -117,6 +117,13 @@ void fromConnectionString() throws MalformedURLException {
         Assertions.assertThat(cfg.getDatabase()).isEqualTo(null);
         Assertions.assertThat(cfg.getWritePrecision()).isEqualTo(WritePrecision.S);
         Assertions.assertThat(cfg.getGzipThreshold()).isEqualTo(1000); // default
+
+        cfg = new ClientConfig.Builder()
+                .build("http://localhost:9999/"
+                        + "?token=my-token&authScheme=my-auth");
+        Assertions.assertThat(cfg.getHost()).isEqualTo("http://localhost:9999/");
+        Assertions.assertThat(cfg.getToken()).isEqualTo("my-token".toCharArray());
+        Assertions.assertThat(cfg.getAuthScheme()).isEqualTo("my-auth");
     }
 
     @Test
@@ -136,7 +143,7 @@ void fromEnv() {
         Assertions.assertThat(cfg.getWritePrecision()).isEqualTo(WritePrecision.NS);
         Assertions.assertThat(cfg.getGzipThreshold()).isEqualTo(1000);
 
-        // simple
+        // basic
         env = Map.of(
                 "INFLUX_HOST", "http://localhost:9999/",
                 "INFLUX_TOKEN", "my-token",
@@ -153,6 +160,18 @@ void fromEnv() {
         Assertions.assertThat(cfg.getWritePrecision()).isEqualTo(WritePrecision.NS);
         Assertions.assertThat(cfg.getGzipThreshold()).isEqualTo(1000);
 
+        // with Edge authentication
+        env = Map.of(
+                "INFLUX_HOST", "http://localhost:9999/",
+                "INFLUX_TOKEN", "my-token",
+                "INFLUX_AUTH_SCHEME", "my-auth"
+        );
+        cfg = new ClientConfig.Builder()
+                .build(env, null);
+        Assertions.assertThat(cfg.getHost()).isEqualTo("http://localhost:9999/");
+        Assertions.assertThat(cfg.getToken()).isEqualTo("my-token".toCharArray());
+        Assertions.assertThat(cfg.getAuthScheme()).isEqualTo("my-auth");
+
         // with write options
         env = Map.of(
                 "INFLUX_HOST", "http://localhost:9999/",
@@ -188,7 +207,7 @@ void fromSystemProperties() {
         Assertions.assertThat(cfg.getWritePrecision()).isEqualTo(WritePrecision.NS);
         Assertions.assertThat(cfg.getGzipThreshold()).isEqualTo(1000);
 
-        // simple
+        // basic
         properties = new Properties();
         properties.put("influx.host", "http://localhost:9999/");
         properties.put("influx.token", "my-token");
@@ -204,6 +223,17 @@ void fromSystemProperties() {
         Assertions.assertThat(cfg.getWritePrecision()).isEqualTo(WritePrecision.NS);
         Assertions.assertThat(cfg.getGzipThreshold()).isEqualTo(1000);
 
+        // with custom auth scheme
+        properties = new Properties();
+        properties.put("influx.host", "http://localhost:9999/");
+        properties.put("influx.token", "my-token");
+        properties.put("influx.authScheme", "my-auth");
+        cfg = new ClientConfig.Builder()
+                .build(new HashMap<>(), properties);
+        Assertions.assertThat(cfg.getHost()).isEqualTo("http://localhost:9999/");
+        Assertions.assertThat(cfg.getToken()).isEqualTo("my-token".toCharArray());
+        Assertions.assertThat(cfg.getAuthScheme()).isEqualTo("my-auth");
+
         // with write options
         properties = new Properties();
         properties.put("influx.host", "http://localhost:9999/");
diff --git a/src/test/java/com/influxdb/v3/client/internal/RestClientTest.java b/src/test/java/com/influxdb/v3/client/internal/RestClientTest.java
@@ -110,6 +110,24 @@ public void authenticationHeader() throws InterruptedException {
         Assertions.assertThat(authorization).isEqualTo("Token my-token");
     }
 
+    @Test
+    public void authenticationHeaderCustomAuthScheme() throws InterruptedException {
+        mockServer.enqueue(createResponse(200));
+
+        restClient = new RestClient(new ClientConfig.Builder()
+                .host(baseURL)
+                .token("my-token".toCharArray())
+                .authScheme("my-auth-scheme")
+                .build());
+
+        restClient.request("ping", HttpMethod.GET, null, null, null);
+
+        RecordedRequest recordedRequest = mockServer.takeRequest();
+
+        String authorization = recordedRequest.getHeader("Authorization");
+        Assertions.assertThat(authorization).isEqualTo("my-auth-scheme my-token");
+    }
+
     @Test
     public void authenticationHeaderNotDefined() throws InterruptedException {
         mockServer.enqueue(createResponse(200));

Original file line number	Diff line number	Diff line change
`@@ -162,7 +162,11 @@ void request(@Nonnull final String path,`
`162`	`162`	`}`
`163`	`163`	`request.header("User-Agent", userAgent);`
`164`	`164`	`if (config.getToken() != null && config.getToken().length > 0) {`
`165`		`- request.header("Authorization", String.format("Token %s", new String(config.getToken())));`
	`165`	`+ String authScheme = config.getAuthScheme();`
	`166`	`+ if (authScheme == null) {`
	`167`	`+ authScheme = "Token";`
	`168`	`+ }`
	`169`	`+ request.header("Authorization", String.format("%s %s", authScheme, new String(config.getToken())));`
`166`	`170`	`}`
`167`	`171`
`168`	`172`	`HttpResponse<String> response;`