feat: move credential handling responsibility to the proxy

stack-info: PR: #9690, branch: igorbernstein2/stack/3

Author: igorbernstein2

bigtable/bigtable-proxy/src/main/java/com/google/cloud/bigtable/examples/proxy/commands/Serve.java

@@ -15,18 +15,22 @@
  */
 package com.google.cloud.bigtable.examples.proxy.commands;
 
+import com.google.auth.Credentials;
+import com.google.auth.oauth2.GoogleCredentials;
 import com.google.bigtable.admin.v2.BigtableInstanceAdminGrpc;
 import com.google.bigtable.admin.v2.BigtableTableAdminGrpc;
 import com.google.bigtable.v2.BigtableGrpc;
 import com.google.cloud.bigtable.examples.proxy.core.ProxyHandler;
 import com.google.cloud.bigtable.examples.proxy.core.Registry;
 import com.google.common.collect.ImmutableMap;
 import com.google.longrunning.OperationsGrpc;
+import io.grpc.CallCredentials;
 import io.grpc.InsecureServerCredentials;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.Server;
 import io.grpc.ServerCallHandler;
+import io.grpc.auth.MoreCallCredentials;
 import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder;
 import java.io.IOException;
 import java.net.InetSocketAddress;
@@ -66,6 +70,7 @@ public class Serve implements Callable<Void> {
 
   ManagedChannel adminChannel = null;
   ManagedChannel dataChannel = null;
+  Credentials credentials = null;
   Server server;
 
   @Override
@@ -93,17 +98,21 @@ void start() throws IOException {
               .disableRetry()
               .build();
     }
+    if (credentials == null) {
+      credentials = GoogleCredentials.getApplicationDefault();
+    }
+    CallCredentials callCredentials = MoreCallCredentials.from(credentials);
 
     Map<String, ServerCallHandler<byte[], byte[]>> serviceMap =
         ImmutableMap.of(
             BigtableGrpc.SERVICE_NAME,
-            new ProxyHandler<>(dataChannel),
+            new ProxyHandler<>(dataChannel, callCredentials),
             BigtableInstanceAdminGrpc.SERVICE_NAME,
-            new ProxyHandler<>(adminChannel),
+            new ProxyHandler<>(adminChannel, callCredentials),
             BigtableTableAdminGrpc.SERVICE_NAME,
-            new ProxyHandler<>(adminChannel),
+            new ProxyHandler<>(adminChannel, callCredentials),
             OperationsGrpc.SERVICE_NAME,
-            new ProxyHandler<>(adminChannel));
+            new ProxyHandler<>(adminChannel, callCredentials));
 
     server =
         NettyServerBuilder.forAddress(

bigtable/bigtable-proxy/src/main/java/com/google/cloud/bigtable/examples/proxy/core/ProxyHandler.java

@@ -15,6 +15,7 @@
  */
 package com.google.cloud.bigtable.examples.proxy.core;
 
+import io.grpc.CallCredentials;
 import io.grpc.CallOptions;
 import io.grpc.Channel;
 import io.grpc.ClientCall;
@@ -24,15 +25,23 @@
 
 /** A factory pairing of an incoming server call to an outgoing client call. */
 public final class ProxyHandler<ReqT, RespT> implements ServerCallHandler<ReqT, RespT> {
+  private final Metadata.Key<String> AUTHORIZATION_KEY =
+      Metadata.Key.of("Authorization", Metadata.ASCII_STRING_MARSHALLER);
+
   private final Channel channel;
+  private final CallCredentials callCredentials;
 
-  public ProxyHandler(Channel channel) {
+  public ProxyHandler(Channel channel, CallCredentials callCredentials) {
     this.channel = channel;
+    this.callCredentials = callCredentials;
   }
 
   @Override
   public ServerCall.Listener<ReqT> startCall(ServerCall<ReqT, RespT> serverCall, Metadata headers) {
-    CallOptions callOptions = CallOptions.DEFAULT;
+    // Strip incoming credentials
+    headers.removeAll(AUTHORIZATION_KEY);
+    // Inject proxy credentials
+    CallOptions callOptions = CallOptions.DEFAULT.withCallCredentials(callCredentials);
 
     ClientCall<ReqT, RespT> clientCall =
         channel.newCall(serverCall.getMethodDescriptor(), callOptions);

bigtable/bigtable-proxy/src/test/java/com/google/cloud/bigtable/examples/proxy/commands/ServeTest.java

@@ -19,6 +19,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
+import com.google.auth.Credentials;
 import com.google.bigtable.admin.v2.BigtableInstanceAdminGrpc;
 import com.google.bigtable.admin.v2.BigtableInstanceAdminGrpc.BigtableInstanceAdminFutureStub;
 import com.google.bigtable.admin.v2.BigtableInstanceAdminGrpc.BigtableInstanceAdminImplBase;
@@ -34,6 +35,7 @@
 import com.google.bigtable.v2.BigtableGrpc.BigtableImplBase;
 import com.google.bigtable.v2.CheckAndMutateRowRequest;
 import com.google.bigtable.v2.CheckAndMutateRowResponse;
+import com.google.common.collect.Lists;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.longrunning.GetOperationRequest;
 import com.google.longrunning.Operation;
@@ -62,6 +64,9 @@
 import io.grpc.testing.GrpcCleanupRule;
 import java.io.IOException;
 import java.net.ServerSocket;
+import java.net.URI;
+import java.util.List;
+import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.BlockingDeque;
 import java.util.concurrent.BlockingQueue;
@@ -93,6 +98,7 @@ public class ServeTest {
   private FakeTableAdminService tableAdminService;
   private OperationService operationService;
   private ManagedChannel fakeServiceChannel;
+  private FakeCredentials fakeCredentials;
 
   // Proxy
   private Serve serve;
@@ -107,6 +113,8 @@ public void setUp() throws IOException {
     tableAdminService = new FakeTableAdminService();
     operationService = new OperationService();
 
+    fakeCredentials = new FakeCredentials();
+
     grpcCleanup.register(
         InProcessServerBuilder.forName(TARGET_SERVER_NAME)
             .intercept(metadataInterceptor)
@@ -122,7 +130,9 @@ public void setUp() throws IOException {
             InProcessChannelBuilder.forName(TARGET_SERVER_NAME).usePlaintext().build());
 
     // Create the proxy
-    serve = createAndStartCommand(fakeServiceChannel);
+    // Inject fakes for upstream calls. For unit tests we want to shim communications to the
+    // bigtable service.
+    serve = createAndStartCommand(fakeServiceChannel, fakeCredentials);
 
     proxyChannel =
         grpcCleanup.register(
@@ -318,11 +328,84 @@ public void onClose(Status status, Metadata trailers) {
     assertThat(clientRecvTrailer.get()).hasValue("trailer", "trailer-value");
   }
 
-  private static Serve createAndStartCommand(ManagedChannel targetChannel) throws IOException {
+  @Test
+  public void testCredentials() throws InterruptedException, ExecutionException, TimeoutException {
+    BigtableFutureStub proxyStub = BigtableGrpc.newFutureStub(proxyChannel);
+
+    CheckAndMutateRowRequest request =
+        CheckAndMutateRowRequest.newBuilder().setTableName("some-table").build();
+    ListenableFuture<CheckAndMutateRowResponse> proxyFuture = proxyStub.checkAndMutateRow(request);
+    StreamObserver<CheckAndMutateRowResponse> serverObserver =
+        dataService
+            .calls
+            .computeIfAbsent(request, (ignored) -> new LinkedBlockingDeque<>())
+            .poll(1, TimeUnit.SECONDS);
+
+    assertWithMessage("Timed out waiting for the proxied RPC on the fake server")
+        .that(serverObserver)
+        .isNotNull();
+
+    serverObserver.onNext(CheckAndMutateRowResponse.newBuilder().setPredicateMatched(true).build());
+    serverObserver.onCompleted();
+    proxyFuture.get(1, TimeUnit.SECONDS);
+
+    assertThat(metadataInterceptor.requestHeaders.poll(1, TimeUnit.SECONDS))
+        .hasValue("authorization", "fake-token");
+  }
+
+  @Test
+  public void testCredentialsClobber()
+      throws InterruptedException, ExecutionException, TimeoutException {
+    BigtableFutureStub proxyStub =
+        BigtableGrpc.newFutureStub(proxyChannel)
+            .withInterceptors(
+                new ClientInterceptor() {
+                  @Override
+                  public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
+                      MethodDescriptor<ReqT, RespT> methodDescriptor,
+                      CallOptions callOptions,
+                      Channel channel) {
+                    return new SimpleForwardingClientCall<ReqT, RespT>(
+                        channel.newCall(methodDescriptor, callOptions)) {
+                      @Override
+                      public void start(Listener<RespT> responseListener, Metadata headers) {
+                        headers.put(
+                            Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER),
+                            "pre-proxied-value");
+                        super.start(responseListener, headers);
+                      }
+                    };
+                  }
+                });
+
+    CheckAndMutateRowRequest request =
+        CheckAndMutateRowRequest.newBuilder().setTableName("some-table").build();
+    ListenableFuture<CheckAndMutateRowResponse> proxyFuture = proxyStub.checkAndMutateRow(request);
+    StreamObserver<CheckAndMutateRowResponse> serverObserver =
+        dataService
+            .calls
+            .computeIfAbsent(request, (ignored) -> new LinkedBlockingDeque<>())
+            .poll(1, TimeUnit.SECONDS);
+
+    assertWithMessage("Timed out waiting for the proxied RPC on the fake server")
+        .that(serverObserver)
+        .isNotNull();
+
+    serverObserver.onNext(CheckAndMutateRowResponse.newBuilder().setPredicateMatched(true).build());
+    serverObserver.onCompleted();
+    proxyFuture.get(1, TimeUnit.SECONDS);
+
+    Metadata serverRequestHeaders = metadataInterceptor.requestHeaders.poll(1, TimeUnit.SECONDS);
+    assertThat(serverRequestHeaders).hasValue("authorization", "fake-token");
+  }
+
+  private static Serve createAndStartCommand(
+      ManagedChannel targetChannel, FakeCredentials targetCredentials) throws IOException {
     for (int i = 10; i >= 0; i--) {
       Serve s = new Serve();
       s.dataChannel = targetChannel;
       s.adminChannel = targetChannel;
+      s.credentials = targetCredentials;
 
       try (ServerSocket serverSocket = new ServerSocket(0)) {
         s.listenPort = serverSocket.getLocalPort();
@@ -419,4 +502,34 @@ public void getOperation(
           .add(responseObserver);
     }
   }
+
+  private static class FakeCredentials extends Credentials {
+    private static final String HEADER_NAME = "authorization";
+    private String fakeValue = "fake-token";
+
+    @Override
+    public String getAuthenticationType() {
+      return "fake";
+    }
+
+    @Override
+    public Map<String, List<String>> getRequestMetadata(URI uri) throws IOException {
+      return Map.of(HEADER_NAME, Lists.newArrayList(fakeValue));
+    }
+
+    @Override
+    public boolean hasRequestMetadata() {
+      return true;
+    }
+
+    @Override
+    public boolean hasRequestMetadataOnly() {
+      return true;
+    }
+
+    @Override
+    public void refresh() throws IOException {
+      // noop
+    }
+  }
 }

bigtable/bigtable-proxy/src/test/java/com/google/cloud/bigtable/examples/proxy/utils/MetadataSubject.java

@@ -5,6 +5,8 @@
 import com.google.common.truth.FailureMetadata;
 import com.google.common.truth.Subject;
 import io.grpc.Metadata;
+import java.util.ArrayList;
+import java.util.Optional;
 import org.jspecify.annotations.Nullable;
 
 public class MetadataSubject extends Subject {
@@ -36,6 +38,17 @@ public void hasValue(String key, String value) {
   }
 
   public <T> void hasValue(Metadata.Key<T> key, T value) {
-    check("get(" + key + ")").that(metadata.get(key)).isEqualTo(value);
+    Iterable<T> actualValues = Optional.ofNullable(metadata.getAll(key)).orElse(new ArrayList<>());
+    check("get(" + key + ")").that(actualValues).containsExactly(value);
+  }
+
+  public void containsValue(String key, String value) {
+    check("get(" + key + ")")
+        .that(metadata.getAll(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER)))
+        .contains(value);
+  }
+
+  public <T> void containsValue(Metadata.Key<T> key, T value) {
+    check("get(" + key + ")").that(metadata.getAll(key)).contains(value);
   }
 }