[CALCITE-6135] BEARER authentication support

Co-authored-by: Aron Meszaros <meszaros.aron.attila@gmail.com>

Author: Richard Antal

bom/build.gradle.kts

@@ -79,6 +79,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

@@ -44,6 +44,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

@@ -137,7 +137,16 @@ public enum BuiltInConnectionProperty implements ConnectionProperty {
    * HTTP Response Timeout (socket timeout) in milliseconds.
    */
   HTTP_RESPONSE_TIMEOUT("http_response_timeout",
-      Type.NUMBER, Timeout.ofMinutes(3).toMilliseconds(), false);
+                        Type.NUMBER, Timeout.ofMinutes(3).toMilliseconds(), false),
+
+  /** Bearer token to use to perform Bearer authentication. */
+  BEARER_TOKEN("bearer_token", Type.STRING, null, false),
+
+  /** File containing bearer tokens to use to perform Bearer authentication. */
+  TOKEN_FILE("token_file", Type.STRING, "", false),
+
+  /** Classname of the BearerTokenProvider. */
+  TOKEN_PROVIDER_CLASS("bearer_token_provider_class", Type.STRING, null, false);
 
   private final String camelName;
   private final Type type;

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

@@ -84,6 +84,14 @@ public interface ConnectionConfig {
   long getHttpConnectionTimeout();
   /** @see BuiltInConnectionProperty#HTTP_RESPONSE_TIMEOUT **/
   long getHttpResponseTimeout();
+  /** @see BuiltInConnectionProperty#TOKEN_FILE */
+  String getTokenFile();
+  /** @see BuiltInConnectionProperty#BEARER_TOKEN */
+  String getBearerToken();
+  /** @see BuiltInConnectionProperty#TOKEN_PROVIDER_CLASS */
+  String getBearerTokenProviderClass();
+
+  ConnectionPropertyValue customPropertyValue(ConnectionProperty property);
 }
 
 // End ConnectionConfig.java

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

@@ -175,6 +175,22 @@ public long getHttpResponseTimeout() {
     return BuiltInConnectionProperty.HTTP_RESPONSE_TIMEOUT.wrap(properties).getLong();
   }
 
+  public String getTokenFile() {
+    return BuiltInConnectionProperty.TOKEN_FILE.wrap(properties).getString();
+  }
+
+  public String getBearerToken() {
+    return BuiltInConnectionProperty.BEARER_TOKEN.wrap(properties).getString();
+  }
+
+  public String getBearerTokenProviderClass() {
+    return BuiltInConnectionProperty.TOKEN_PROVIDER_CLASS.wrap(properties).getString();
+  }
+
+  public ConnectionPropertyValue customPropertyValue(ConnectionProperty property) {
+    return property.wrap(properties);
+  }
+
   /** Converts a {@link Properties} object containing (name, value)
    * pairs into a map whose keys are
    * {@link org.apache.calcite.avatica.InternalProperty} objects.
@@ -204,7 +220,7 @@ public static Map<ConnectionProperty, String> parse(Properties properties,
   }
 
   /** The combination of a property definition and a map of property values. */
-  public static class PropEnv {
+  public static class PropEnv implements ConnectionPropertyValue {
     final Map<? extends ConnectionProperty, String> map;
     private final ConnectionProperty property;
 

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

@@ -40,7 +40,7 @@ public interface ConnectionProperty {
 
   /** Wraps this property with a properties object from which its value can be
    * obtained when needed. */
-  ConnectionConfigImpl.PropEnv wrap(Properties properties);
+  ConnectionPropertyValue wrap(Properties properties);
 
   /** Whether the property is mandatory. */
   boolean required();

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

@@ -0,0 +1,113 @@
+/*
+ * 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;
+
+public interface ConnectionPropertyValue {
+  /**
+   * Returns the string value of this property, or null if not specified and
+   * no default.
+   */
+  String getString();
+
+  /**
+   * Returns the string value of this property, or null if not specified and
+   * no default.
+   */
+  String getString(String defaultValue);
+
+  /**
+   * Returns the int value of this property. Throws if not set and no
+   * default.
+   */
+  int getInt();
+
+  /**
+   * Returns the int value of this property. Throws if not set and no
+   * default.
+   */
+  int getInt(Number defaultValue);
+
+  /**
+   * Returns the long value of this property. Throws if not set and no
+   * default.
+   */
+  long getLong();
+
+  /**
+   * Returns the long value of this property. Throws if not set and no
+   * default.
+   */
+  long getLong(Number defaultValue);
+
+  /**
+   * Returns the double value of this property. Throws if not set and no
+   * default.
+   */
+  double getDouble();
+
+  /**
+   * Returns the double value of this property. Throws if not set and no
+   * default.
+   */
+  double getDouble(Number defaultValue);
+
+  /**
+   * Returns the boolean value of this property. Throws if not set and no
+   * default.
+   */
+  boolean getBoolean();
+
+  /**
+   * Returns the boolean value of this property. Throws if not set and no
+   * default.
+   */
+  boolean getBoolean(boolean defaultValue);
+
+  /**
+   * Returns the enum value of this property. Throws if not set and no
+   * default.
+   */
+  <E extends Enum<E>> E getEnum(Class<E> enumClass);
+
+  /**
+   * Returns the enum value of this property. Throws if not set and no
+   * default.
+   */
+  <E extends Enum<E>> E getEnum(Class<E> enumClass, E defaultValue);
+
+  /**
+   * Returns an instance of a plugin.
+   *
+   * <p>Throws if not set and no default.
+   * Also throws if the class does not implement the required interface,
+   * or if it does not have a public default constructor or an public static
+   * field called {@code #INSTANCE}.
+   */
+  <T> T getPlugin(Class<T> pluginClass, T defaultInstance);
+
+  /**
+   * Returns an instance of a plugin, using a given class name if none is
+   * set.
+   *
+   * <p>Throws if not set and no default.
+   * Also throws if the class does not implement the required interface,
+   * or if it does not have a public default constructor or an public static
+   * field called {@code #INSTANCE}.
+   */
+  <T> T getPlugin(Class<T> pluginClass, String defaultClassName,
+                  T defaultInstance);
+}

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

@@ -22,6 +22,7 @@
 import org.apache.hc.client5.http.SystemDefaultDnsResolver;
 import org.apache.hc.client5.http.auth.AuthSchemeFactory;
 import org.apache.hc.client5.http.auth.AuthScope;
+import org.apache.hc.client5.http.auth.BearerToken;
 import org.apache.hc.client5.http.auth.Credentials;
 import org.apache.hc.client5.http.auth.CredentialsProvider;
 import org.apache.hc.client5.http.auth.StandardAuthScheme;
@@ -31,6 +32,7 @@
 import org.apache.hc.client5.http.impl.auth.BasicAuthCache;
 import org.apache.hc.client5.http.impl.auth.BasicCredentialsProvider;
 import org.apache.hc.client5.http.impl.auth.BasicSchemeFactory;
+import org.apache.hc.client5.http.impl.auth.BearerSchemeFactory;
 import org.apache.hc.client5.http.impl.auth.DigestSchemeFactory;
 import org.apache.hc.client5.http.impl.classic.CloseableHttpClient;
 import org.apache.hc.client5.http.impl.classic.HttpClientBuilder;
@@ -68,7 +70,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
@@ -152,6 +154,9 @@ private RequestConfig createRequestConfig() {
       if (authRegistry.lookup(StandardAuthScheme.SPNEGO) != null) {
         preferredSchemes.add(StandardAuthScheme.SPNEGO);
       }
+      if (authRegistry.lookup(StandardAuthScheme.BEARER) != null) {
+        preferredSchemes.add(StandardAuthScheme.BEARER);
+      }
       requestConfigBuilder.setTargetPreferredAuthSchemes(preferredSchemes);
       requestConfigBuilder.setProxyPreferredAuthSchemes(preferredSchemes);
     }
@@ -258,6 +263,20 @@ ClassicHttpResponse executeOpen(HttpHost httpHost, HttpPost post, HttpClientCont
     context.setRequestConfig(createRequestConfig());
   }
 
+  @Override public void setTokenProvider(String username, BearerTokenProvider tokenProvider) {
+    this.credentialsProvider = new BasicCredentialsProvider();
+    ((BasicCredentialsProvider) this.credentialsProvider)
+        .setCredentials(anyAuthScope, new BearerToken(tokenProvider.obtain(username)));
+
+    this.authRegistry = RegistryBuilder.<AuthSchemeFactory>create()
+        .register(StandardAuthScheme.BEARER,
+            new BearerSchemeFactory())
+        .build();
+    context.setCredentialsProvider(credentialsProvider);
+    context.setAuthSchemeRegistry(authRegistry);
+    context.setRequestConfig(createRequestConfig());
+  }
+
   /**
    * A credentials implementation which returns null.
    */

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

@@ -131,6 +131,24 @@ public static AvaticaHttpClientFactoryImpl getInstance() {
       LOG.debug("{} is not capable of kerberos authentication.", authType);
     }
 
+    if (client instanceof BearerAuthenticateable) {
+      if (AuthenticationType.BEARER == authType) {
+        try {
+          BearerTokenProvider tokenProvider =
+                  BearerTokenProviderFactory.getBearerTokenProvider(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);
     }