feat: Test proxy support SSL backend (#2381)

jackdingilian · Patrick Wrobel · web-flow · commit 3cbf4abe79d6 · 2024-10-22T12:04:41.000-04:00
Change-Id: I622b1cca3b02d176beaef8a21dce8bab6b16a937

Co-authored-by: Patrick Wrobel &lt;pwrobel@google.com&gt;
diff --git a/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java b/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java
@@ -26,7 +26,8 @@
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.rpc.ApiException;
 import com.google.api.gax.rpc.ServerStream;
-import com.google.auth.oauth2.GoogleCredentials;
+import com.google.auth.oauth2.AccessToken;
+import com.google.auth.oauth2.OAuth2Credentials;
 import com.google.auto.value.AutoValue;
 import com.google.bigtable.v2.Column;
 import com.google.bigtable.v2.Family;
@@ -59,9 +60,6 @@
 import java.io.ByteArrayInputStream;
 import java.io.Closeable;
 import java.io.IOException;
-import java.nio.file.Files;
-import java.nio.file.Path;
-import java.nio.file.Paths;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -72,7 +70,6 @@
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
-import javax.annotation.Nullable;
 import org.threeten.bp.Duration;
 
 /** Java implementation of the CBT test proxy. Used to test the Java CBT client. */
@@ -95,50 +92,13 @@ static CbtClient create(BigtableDataSettings settings, BigtableDataClient dataCl
 
   private static final Logger logger = Logger.getLogger(CbtTestProxy.class.getName());
 
-  private CbtTestProxy(
-      boolean encrypted,
-      @Nullable String rootCerts,
-      @Nullable String sslTarget,
-      @Nullable String credential) {
-    this.encrypted = encrypted;
-    this.rootCerts = rootCerts;
-    this.sslTarget = sslTarget;
-    this.credential = credential;
+  private CbtTestProxy() {
     this.idClientMap = new ConcurrentHashMap<>();
   }
 
-  /**
-   * Factory method to return a proxy instance that interacts with server unencrypted and
-   * unauthenticated.
-   */
-  public static CbtTestProxy createUnencrypted() {
-    return new CbtTestProxy(false, null, null, null);
-  }
-
-  /**
-   * Factory method to return a proxy instance that interacts with server encrypted. Default
-   * authority and public certificates are used if null values are passed in.
-   *
-   * @param rootCertsPemPath The path to a root certificate PEM file
-   * @param sslTarget The override of SSL target name
-   * @param credentialJsonPath The path to a credential JSON file
-   */
-  public static CbtTestProxy createEncrypted(
-      @Nullable String rootCertsPemPath,
-      @Nullable String sslTarget,
-      @Nullable String credentialJsonPath)
-      throws IOException {
-    String tmpRootCerts = null, tmpCredential = null;
-    if (rootCertsPemPath != null) {
-      Path file = Paths.get(rootCertsPemPath);
-      tmpRootCerts = new String(Files.readAllBytes(file), UTF_8);
-    }
-    if (credentialJsonPath != null) {
-      Path file = Paths.get(credentialJsonPath);
-      tmpCredential = new String(Files.readAllBytes(file), UTF_8);
-    }
-
-    return new CbtTestProxy(true, tmpRootCerts, sslTarget, tmpCredential);
+  /** Factory method to return a proxy instance. */
+  public static CbtTestProxy create() {
+    return new CbtTestProxy();
   }
 
   /**
@@ -196,15 +156,21 @@ public synchronized void createClient(
     Preconditions.checkArgument(!request.getProjectId().isEmpty(), "project id must be provided");
     Preconditions.checkArgument(!request.getInstanceId().isEmpty(), "instance id must be provided");
     Preconditions.checkArgument(!request.getDataTarget().isEmpty(), "data target must be provided");
+    Preconditions.checkArgument(
+        !request.getSecurityOptions().getUseSsl()
+            || !request.getSecurityOptions().getSslRootCertsPemBytes().isEmpty(),
+        "security_options.ssl_root_certs_pem must be provided if security_options.use_ssl is true");
 
-    if (idClientMap.contains(request.getClientId())) {
+    if (idClientMap.containsKey(request.getClientId())) {
       responseObserver.onError(
           Status.ALREADY_EXISTS
               .withDescription("Client " + request.getClientId() + " already exists.")
               .asException());
       return;
     }
 
+    // setRefreshingChannel is needed for now.
+    @SuppressWarnings("deprecation")
     BigtableDataSettings.Builder settingsBuilder =
         BigtableDataSettings.newBuilder()
             // Disable channel refreshing when not using the real server
@@ -213,9 +179,6 @@ public synchronized void createClient(
             .setInstanceId(request.getInstanceId())
             .setAppProfileId(request.getAppProfileId());
 
-    settingsBuilder.stubSettings().setEnableRoutingCookie(false);
-    settingsBuilder.stubSettings().setEnableRetryInfo(false);
-
     if (request.hasPerOperationTimeout()) {
       Duration newTimeout = Duration.ofMillis(Durations.toMillis(request.getPerOperationTimeout()));
       settingsBuilder = overrideTimeoutSetting(newTimeout, settingsBuilder);
@@ -249,8 +212,13 @@ public synchronized void createClient(
         settingsBuilder
             .stubSettings()
             .setEndpoint(request.getDataTarget())
-            .setTransportChannelProvider(getTransportChannel())
-            .setCredentialsProvider(getCredentialsProvider());
+            .setTransportChannelProvider(
+                getTransportChannel(
+                    request.getSecurityOptions().getUseSsl(),
+                    request.getSecurityOptions().getSslRootCertsPem(),
+                    request.getSecurityOptions().getSslEndpointOverride()))
+            .setCredentialsProvider(
+                getCredentialsProvider(request.getSecurityOptions().getAccessToken()));
       }
       BigtableDataSettings settings = settingsBuilder.build();
       BigtableDataClient client = BigtableDataClient.create(settings);
@@ -780,52 +748,60 @@ private static String extractTableIdFromTableName(String fullTableName)
     return matcher.group(3);
   }
 
-  private InstantiatingGrpcChannelProvider getTransportChannel() throws IOException {
+  @SuppressWarnings("rawtypes")
+  private InstantiatingGrpcChannelProvider getTransportChannel(
+      boolean encrypted, String rootCertsPem, String sslTarget) {
     if (!encrypted) {
       return EnhancedBigtableStubSettings.defaultGrpcTransportProviderBuilder()
           .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
           .build();
     }
 
-    if (rootCerts == null) {
-      return EnhancedBigtableStubSettings.defaultGrpcTransportProviderBuilder().build();
+    final SslContext sslContext;
+    if (rootCertsPem.isEmpty()) {
+      sslContext = null;
+    } else {
+      try {
+        sslContext =
+            GrpcSslContexts.forClient()
+                .trustManager(new ByteArrayInputStream(rootCertsPem.getBytes(UTF_8)))
+                .build();
+      } catch (IOException e) {
+        throw new IllegalArgumentException(e);
+      }
     }
 
-    final SslContext secureContext =
-        GrpcSslContexts.forClient()
-            .trustManager(new ByteArrayInputStream(rootCerts.getBytes(UTF_8)))
-            .build();
     return EnhancedBigtableStubSettings.defaultGrpcTransportProviderBuilder()
         .setChannelConfigurator(
             new ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder>() {
               @Override
               public ManagedChannelBuilder apply(ManagedChannelBuilder input) {
                 NettyChannelBuilder channelBuilder = (NettyChannelBuilder) input;
-                channelBuilder.sslContext(secureContext).overrideAuthority(sslTarget);
+
+                if (sslContext != null) {
+                  channelBuilder.sslContext(sslContext);
+                }
+
+                if (!sslTarget.isEmpty()) {
+                  channelBuilder.overrideAuthority(sslTarget);
+                }
+
                 return channelBuilder;
               }
             })
         .build();
   }
 
-  private CredentialsProvider getCredentialsProvider() throws IOException {
-    if (credential == null) {
+  private CredentialsProvider getCredentialsProvider(String accessToken) {
+    if (accessToken.isEmpty()) {
       return NoCredentialsProvider.create();
     }
 
-    final GoogleCredentials creds =
-        GoogleCredentials.fromStream(new ByteArrayInputStream(credential.getBytes(UTF_8)));
-
-    return FixedCredentialsProvider.create(creds);
+    return FixedCredentialsProvider.create(
+        OAuth2Credentials.create(new AccessToken(accessToken, null)));
   }
 
   private final ConcurrentHashMap<String, CbtClient> idClientMap;
-  private final boolean encrypted;
-
-  // Parameters that may be needed when "encrypted" is true.
-  private final String rootCerts;
-  private final String sslTarget;
-  private final String credential;
 
   private static final Pattern tablePattern =
       Pattern.compile("projects/([^/]+)/instances/([^/]+)/tables/([^/]+)");
diff --git a/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxyMain.java b/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxyMain.java
@@ -32,19 +32,7 @@ public static void main(String[] args) throws InterruptedException, IOException
       throw new IllegalArgumentException(String.format("Port %d is not > 0.", port));
     }
 
-    CbtTestProxy cbtTestProxy;
-
-    // If encryption is specified
-    boolean encrypted = Boolean.getBoolean("encrypted");
-    if (encrypted) {
-      String rootCertsPemPath = System.getProperty("root.certs.pem.path");
-      String sslTarget = System.getProperty("ssl.target");
-      String credentialJsonPath = System.getProperty("credential.json.path");
-      cbtTestProxy = CbtTestProxy.createEncrypted(rootCertsPemPath, sslTarget, credentialJsonPath);
-    } else {
-      cbtTestProxy = CbtTestProxy.createUnencrypted();
-    }
-
+    CbtTestProxy cbtTestProxy = CbtTestProxy.create();
     logger.info(String.format("Test proxy starting on %d", port));
     ServerBuilder.forPort(port).addService(cbtTestProxy).build().start().awaitTermination();
   }
diff --git a/test-proxy/src/main/proto/test_proxy.proto b/test-proxy/src/main/proto/test_proxy.proto
@@ -38,6 +38,27 @@ enum OptionalFeatureConfig {
 
 // Request to test proxy service to create a client object.
 message CreateClientRequest {
+  message SecurityOptions {
+    // Access token to use for client credentials. If empty, the client will not
+    // use any call credentials. Certain implementations may require `use_ssl`
+    // to be set when using this.
+    string access_token = 1;
+
+    // Whether to use SSL channel credentials when connecting to the data
+    // endpoint.
+    bool use_ssl = 2;
+
+    // If using SSL channel credentials, override the SSL endpoint to match the
+    // host that is specified in the backend's certificate. Also sets the
+    // client's authority header value.
+    string ssl_endpoint_override = 3;
+
+    // PEM encoding of the server root certificates. If not set, the default
+    // root certs will be used instead. The default can be overridden via the
+    // GRPC_DEFAULT_SSL_ROOTS_FILE_PATH env var.
+    string ssl_root_certs_pem = 4;
+  }
+
   // A unique ID associated with the client object to be created.
   string client_id = 1;
 
@@ -66,6 +87,17 @@ message CreateClientRequest {
   // Optional config that dictates how the optional features should be enabled
   // during the client creation. Please check the enum type's docstring above.
   OptionalFeatureConfig optional_feature_config = 7;
+
+  // Options to allow connecting to backends with channel and/or call
+  // credentials. This is needed internally by Cloud Bigtable's own testing
+  // frameworks.It is not necessary to support these fields for client
+  // conformance testing.
+  //
+  // WARNING: this allows the proxy to connect to a real production
+  // CBT backend with the right options, however, the proxy itself is insecure
+  // so it is not recommended to use it with real credentials or outside testing
+  // contexts.
+  SecurityOptions security_options = 8;
 }
 
 // Response from test proxy service for CreateClientRequest.
