[5.4.0] Add integration tests for --experimental_credential_helper. (bazelbuild#16880)

Closes bazelbuild#15996.

PiperOrigin-RevId: 463851754
Change-Id: I13b851454f0b8e6550b2335caa2d802312931929

Author: tjgq

src/test/shell/bazel/remote/remote_execution_test.sh

@@ -71,6 +71,70 @@ else
   declare -r EXE_EXT=""
 fi
 
+function setup_credential_helper() {
+  cat > "${TEST_TMPDIR}/credhelper" <<'EOF'
+#!/usr/bin/env python3
+print("""{"headers":{"Authorization":["Bearer secret_token"]}}""")
+EOF
+  chmod +x "${TEST_TMPDIR}/credhelper"
+}
+
+function test_credential_helper_remote_cache() {
+  setup_credential_helper
+
+  mkdir -p a
+
+  cat > a/BUILD <<'EOF'
+genrule(
+  name = "a",
+  outs = ["a.txt"],
+  cmd = "touch $(OUTS)",
+)
+EOF
+
+  stop_worker
+  start_worker --expected_authorization_token=secret_token
+
+  bazel build \
+      --remote_cache=grpc://localhost:${worker_port} \
+      //a:a >& $TEST_log && fail "Build without credentials should have failed"
+  expect_log "Failed to query remote execution capabilities"
+
+  bazel build \
+      --remote_cache=grpc://localhost:${worker_port} \
+      --experimental_credential_helper="${TEST_TMPDIR}/credhelper" \
+      //a:a >& $TEST_log || fail "Build with credentials should have succeeded"
+}
+
+function test_credential_helper_remote_execution() {
+  setup_credential_helper
+
+  mkdir -p a
+
+  cat > a/BUILD <<'EOF'
+genrule(
+  name = "a",
+  outs = ["a.txt"],
+  cmd = "touch $(OUTS)",
+)
+EOF
+
+  stop_worker
+  start_worker --expected_authorization_token=secret_token
+
+  bazel build \
+      --spawn_strategy=remote \
+      --remote_executor=grpc://localhost:${worker_port} \
+      //a:a >& $TEST_log && fail "Build without credentials should have failed"
+  expect_log "Failed to query remote execution capabilities"
+
+  bazel build \
+      --spawn_strategy=remote \
+      --remote_executor=grpc://localhost:${worker_port} \
+      --experimental_credential_helper="${TEST_TMPDIR}/credhelper" \
+      //a:a >& $TEST_log || fail "Build with credentials should have succeeded"
+}
+
 function test_remote_grpc_cache_with_protocol() {
   # Test that if 'grpc' is provided as a scheme for --remote_cache flag, remote cache works.
   mkdir -p a

src/tools/remote/src/main/java/com/google/devtools/build/remote/worker/RemoteWorker.java

@@ -49,9 +49,16 @@
 import com.google.devtools.build.remote.worker.http.HttpCacheServerInitializer;
 import com.google.devtools.common.options.OptionsParser;
 import com.google.devtools.common.options.OptionsParsingException;
+import io.grpc.Context;
+import io.grpc.Contexts;
+import io.grpc.Metadata;
 import io.grpc.Server;
+import io.grpc.ServerCall;
+import io.grpc.ServerCall.Listener;
+import io.grpc.ServerCallHandler;
 import io.grpc.ServerInterceptor;
 import io.grpc.ServerInterceptors;
+import io.grpc.Status;
 import io.grpc.netty.GrpcSslContexts;
 import io.grpc.netty.NettyServerBuilder;
 import io.netty.bootstrap.ServerBootstrap;
@@ -71,6 +78,9 @@
 import java.io.OutputStreamWriter;
 import java.io.Writer;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executors;
 import java.util.logging.Level;
@@ -107,6 +117,39 @@ static FileSystem getFileSystem() {
     return new JavaIoFileSystem(hashFunction);
   }
 
+  /** A {@link ServerInterceptor} that rejects requests unless an authorization token is present. */
+  private static class AuthorizationTokenInterceptor implements ServerInterceptor {
+    private static final Metadata.Key<String> AUTHORIZATION_HEADER_KEY =
+        Metadata.Key.of("Authorization", Metadata.ASCII_STRING_MARSHALLER);
+
+    private static final String BEARER_PREFIX = "Bearer ";
+
+    private final String expectedToken;
+
+    AuthorizationTokenInterceptor(String expectedToken) {
+      this.expectedToken = expectedToken;
+    }
+
+    private Optional<String> getTokenFromMetadata(Metadata headers) {
+      String val = headers.get(AUTHORIZATION_HEADER_KEY);
+      if (val != null && val.startsWith(BEARER_PREFIX)) {
+        return Optional.of(val.substring(BEARER_PREFIX.length()));
+      }
+      return Optional.empty();
+    }
+
+    @Override
+    public <ReqT, RespT> Listener<ReqT> interceptCall(
+        ServerCall<ReqT, RespT> call, Metadata headers, ServerCallHandler<ReqT, RespT> next) {
+      Optional<String> actualToken = getTokenFromMetadata(headers);
+      if (!expectedToken.equals(actualToken.get())) {
+        call.close(Status.PERMISSION_DENIED, new Metadata());
+        return new ServerCall.Listener<ReqT>() {};
+      }
+      return Contexts.interceptCall(Context.current(), call, headers, next);
+    }
+  }
+
   public RemoteWorker(
       FileSystem fs,
       RemoteWorkerOptions workerOptions,
@@ -149,20 +192,25 @@ public RemoteWorker(
   }
 
   public Server startServer() throws IOException {
-    ServerInterceptor headersInterceptor = new TracingMetadataUtils.ServerHeadersInterceptor();
+    List<ServerInterceptor> interceptors = new ArrayList<>();
+    interceptors.add(new TracingMetadataUtils.ServerHeadersInterceptor());
+    if (workerOptions.expectedAuthorizationToken != null) {
+      interceptors.add(new AuthorizationTokenInterceptor(workerOptions.expectedAuthorizationToken));
+    }
+
     NettyServerBuilder b =
         NettyServerBuilder.forPort(workerOptions.listenPort)
-            .addService(ServerInterceptors.intercept(actionCacheServer, headersInterceptor))
-            .addService(ServerInterceptors.intercept(bsServer, headersInterceptor))
-            .addService(ServerInterceptors.intercept(casServer, headersInterceptor))
-            .addService(ServerInterceptors.intercept(capabilitiesServer, headersInterceptor));
+            .addService(ServerInterceptors.intercept(actionCacheServer, interceptors))
+            .addService(ServerInterceptors.intercept(bsServer, interceptors))
+            .addService(ServerInterceptors.intercept(casServer, interceptors))
+            .addService(ServerInterceptors.intercept(capabilitiesServer, interceptors));
 
     if (workerOptions.tlsCertificate != null) {
       b.sslContext(getSslContextBuilder(workerOptions).build());
     }
 
     if (execServer != null) {
-      b.addService(ServerInterceptors.intercept(execServer, headersInterceptor));
+      b.addService(ServerInterceptors.intercept(execServer, interceptors));
     } else {
       logger.atInfo().log("Execution disabled, only serving cache requests");
     }

src/tools/remote/src/main/java/com/google/devtools/build/remote/worker/RemoteWorkerOptions.java

@@ -176,6 +176,16 @@ public class RemoteWorkerOptions extends OptionsBase {
               + "requires client authentication (aka mTLS).")
   public String tlsCaCertificate;
 
+  @Option(
+      name = "expected_authorization_token",
+      defaultValue = "null",
+      documentationCategory = OptionDocumentationCategory.UNCATEGORIZED,
+      effectTags = {OptionEffectTag.UNKNOWN},
+      help =
+          "The authorization token expected to be present in every request. This is useful for"
+              + " testing only.")
+  public String expectedAuthorizationToken;
+
   private static final int MAX_JOBS = 16384;
 
   /**