[CALCITE-6135] BEARER authentication support

Author: Aarchy

bom/build.gradle.kts

@@ -78,6 +78,7 @@ dependencies {
         apiv("org.ow2.asm:asm-tree", "asm")
         apiv("org.ow2.asm:asm-util", "asm")
         apiv("org.slf4j:slf4j-api", "slf4j")
+        apiv("commons-io:commons-io")
         // The log4j2 binding should be a runtime dependency but given that
         // some modules shade this dependency we need to keep it as api
         apiv("org.apache.logging.log4j:log4j-slf4j-impl", "log4j2")

core/build.gradle.kts

@@ -42,6 +42,7 @@ dependencies {
     testImplementation("org.mockito:mockito-core")
     testImplementation("org.mockito:mockito-inline")
     testImplementation("org.hamcrest:hamcrest-core")
+    testImplementation("commons-io:commons-io")
     testRuntimeOnly("org.apache.logging.log4j:log4j-slf4j-impl")
 }
 

core/src/main/java/org/apache/calcite/avatica/BuiltInConnectionProperty.java

@@ -127,7 +127,13 @@ public enum BuiltInConnectionProperty implements ConnectionProperty {
    * HTTP Connection Timeout in milliseconds.
    */
   HTTP_CONNECTION_TIMEOUT("http_connection_timeout",
-      Type.NUMBER, Timeout.ofMinutes(3).toMilliseconds(), false);
+      Type.NUMBER, Timeout.ofMinutes(3).toMilliseconds(), false),
+
+  /** File containing bearer tokens to use to perform Bearer authentication. */
+  TOKEN_FILE("tokenfile", Type.STRING, "", false),
+
+  /** Bearer token to use to perform Bearer authentication. */
+  BEARER_TOKEN("bearertoken", Type.STRING, "", false);
 
   private final String camelName;
   private final Type type;

core/src/main/java/org/apache/calcite/avatica/ConnectionConfig.java

@@ -81,6 +81,10 @@ public interface ConnectionConfig {
   long getLBConnectionFailoverSleepTime();
   /** @see BuiltInConnectionProperty#HTTP_CONNECTION_TIMEOUT **/
   long getHttpConnectionTimeout();
+  /** @see BuiltInConnectionProperty#TOKEN_FILE */
+  String tokenFile();
+  /** @see BuiltInConnectionProperty#BEARER_TOKEN */
+  String bearerToken();
 }
 
 // End ConnectionConfig.java

core/src/main/java/org/apache/calcite/avatica/ConnectionConfigImpl.java

@@ -169,6 +169,14 @@ public long getHttpConnectionTimeout() {
     return BuiltInConnectionProperty.HTTP_CONNECTION_TIMEOUT.wrap(properties).getLong();
   }
 
+  public String tokenFile() {
+    return BuiltInConnectionProperty.TOKEN_FILE.wrap(properties).getString();
+  }
+
+  public String bearerToken() {
+    return BuiltInConnectionProperty.BEARER_TOKEN.wrap(properties).getString();
+  }
+
   /** Converts a {@link Properties} object containing (name, value)
    * pairs into a map whose keys are
    * {@link org.apache.calcite.avatica.InternalProperty} objects.

core/src/main/java/org/apache/calcite/avatica/remote/AuthenticationType.java

@@ -24,6 +24,7 @@ public enum AuthenticationType {
   BASIC,
   DIGEST,
   SPNEGO,
+  BEARER,
   CUSTOM;
 }
 

core/src/main/java/org/apache/calcite/avatica/remote/AvaticaCommonsHttpClientImpl.java

@@ -60,6 +60,9 @@
 import java.net.URISyntaxException;
 import java.net.URL;
 import java.security.Principal;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 
@@ -68,7 +71,7 @@
  * sent and received across the wire.
  */
 public class AvaticaCommonsHttpClientImpl implements AvaticaHttpClient, HttpClientPoolConfigurable,
-    UsernamePasswordAuthenticateable, GSSAuthenticateable {
+    UsernamePasswordAuthenticateable, GSSAuthenticateable, BearerAuthenticateable {
   private static final Logger LOG = LoggerFactory.getLogger(AvaticaCommonsHttpClientImpl.class);
 
   // SPNEGO specific settings
@@ -91,6 +94,10 @@ public class AvaticaCommonsHttpClientImpl implements AvaticaHttpClient, HttpClie
   protected CredentialsProvider credentialsProvider = null;
   protected Lookup<AuthSchemeFactory> authRegistry = null;
   protected Object userToken;
+  private static final List<String> AVATICA_SCHEME_PRIORITY =
+      Collections.unmodifiableList(Arrays.asList(StandardAuthScheme.BASIC,
+          StandardAuthScheme.DIGEST, StandardAuthScheme.SPNEGO, StandardAuthScheme.NTLM,
+          StandardAuthScheme.KERBEROS, "Bearer"));
 
   public AvaticaCommonsHttpClientImpl(URL url) {
     this.uri = toURI(Objects.requireNonNull(url));
@@ -104,6 +111,7 @@ protected void initializeClient(PoolingHttpClientConnectionManager pool,
     RequestConfig.Builder requestConfigBuilder = RequestConfig.custom();
     RequestConfig requestConfig = requestConfigBuilder
         .setConnectTimeout(config.getHttpConnectionTimeout(), TimeUnit.MILLISECONDS)
+        .setTargetPreferredAuthSchemes(AVATICA_SCHEME_PRIORITY)
         .build();
     HttpClientBuilder httpClientBuilder = HttpClients.custom().setConnectionManager(pool)
         .setDefaultRequestConfig(requestConfig);
@@ -206,6 +214,17 @@ CloseableHttpResponse execute(HttpPost post, HttpClientContext context)
     }
   }
 
+  @Override public void setTokenProvider(String username, BearerTokenProvider tokenProvider) {
+    this.credentialsProvider = new BasicCredentialsProvider();
+    ((BasicCredentialsProvider) this.credentialsProvider)
+        .setCredentials(anyAuthScope, new BearerCredentials(username, tokenProvider));
+
+    this.authRegistry = RegistryBuilder.<AuthSchemeFactory>create()
+        .register("Bearer",
+            new BearerSchemeFactory())
+        .build();
+  }
+
   /**
    * A credentials implementation which returns null.
    */

core/src/main/java/org/apache/calcite/avatica/remote/AvaticaHttpClientFactoryImpl.java

@@ -130,6 +130,36 @@ public static AvaticaHttpClientFactoryImpl getInstance() {
       LOG.debug("{} is not capable of kerberos authentication.", authType);
     }
 
+    if (client instanceof BearerAuthenticateable) {
+      if (AuthenticationType.BEARER == authType) {
+        if (config.bearerToken() == null && config.tokenFile() == null
+            || config.bearerToken() != null && config.tokenFile() != null) {
+          LOG.debug("Failed to initialize bearer authentication:"
+              + "either the tokenfile or bearertoken must be set");
+        } else {
+          BearerTokenProvider tokenProvider;
+          if (config.bearerToken() != null) {
+            tokenProvider = new SingleBearerTokenProvider();
+          } else {
+            tokenProvider = new FileBearerTokenProvider();
+          }
+
+          try {
+            tokenProvider.init(config);
+            String username = config.avaticaUser();
+            if (null == username) {
+              username = System.getProperty("user.name");
+            }
+            ((BearerAuthenticateable) client).setTokenProvider(username, tokenProvider);
+          } catch (java.io.IOException e) {
+            LOG.debug("Failed to initialize bearer authentication");
+          }
+        }
+      }
+    } else {
+      LOG.debug("{} is not capable of bearer authentication.", authType);
+    }
+
     if (null != kerberosUtil) {
       client = new DoAsAvaticaHttpClient(client, kerberosUtil);
     }

core/src/main/java/org/apache/calcite/avatica/remote/BearerAuthenticateable.java

@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to you under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.calcite.avatica.remote;
+
+public interface BearerAuthenticateable {
+
+  void setTokenProvider(String username, BearerTokenProvider tokenProvider);
+}

core/src/main/java/org/apache/calcite/avatica/remote/BearerCredentials.java

@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to you under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.calcite.avatica.remote;
+
+import org.apache.hc.client5.http.auth.Credentials;
+import org.apache.hc.core5.annotation.Contract;
+import org.apache.hc.core5.annotation.ThreadingBehavior;
+import org.apache.hc.core5.util.Args;
+
+import java.io.Serializable;
+import java.security.Principal;
+
+@Contract(threading = ThreadingBehavior.IMMUTABLE)
+public class BearerCredentials implements Credentials, Serializable {
+
+  private final BearerTokenProvider tokenProvider;
+
+  private final String userName;
+
+  public BearerCredentials(final String userName, final BearerTokenProvider tokenProvider) {
+    Args.notNull(userName, "userName");
+    Args.notNull(tokenProvider, "tokenProvider");
+    this.tokenProvider = tokenProvider;
+    this.userName = userName;
+  }
+
+  public String getToken() {
+    return tokenProvider.obtain(userName);
+  }
+
+  @Override
+  public Principal getUserPrincipal() {
+    return null;
+  }
+
+  @Override
+  public char[] getPassword() {
+    return null;
+  }
+}