Skip to content

Commit

Permalink
Add integration tests for --experimental_credential_helper.
Browse files Browse the repository at this point in the history
  • Loading branch information
tjgq committed Jul 28, 2022
1 parent 1d4cecf commit 4b97d8b
Show file tree
Hide file tree
Showing 3 changed files with 133 additions and 6 deletions.
64 changes: 64 additions & 0 deletions src/test/shell/bazel/remote/remote_execution_test.sh
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,70 @@ function has_utf8_locale() {
[[ "${charmap}" == "UTF-8" ]]
}

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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import build.bazel.remote.execution.v2.CapabilitiesGrpc.CapabilitiesImplBase;
import build.bazel.remote.execution.v2.ContentAddressableStorageGrpc.ContentAddressableStorageImplBase;
import build.bazel.remote.execution.v2.ExecutionGrpc.ExecutionImplBase;
import build.bazel.remote.execution.v2.RequestMetadata;
import com.google.bytestream.ByteStreamGrpc.ByteStreamImplBase;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
Expand Down Expand Up @@ -49,9 +50,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;
Expand All @@ -71,10 +79,15 @@
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;
import java.util.logging.Logger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
* Implements a remote worker that accepts work items as protobufs. The server implementation is
Expand Down Expand Up @@ -107,6 +120,41 @@ 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);
logger.atInfo().log("headers: %s, token: %s", headers, actualToken);
if (!expectedToken.equals(actualToken.get())) {
logger.atInfo().log("permission denied");
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,
Expand Down Expand Up @@ -149,20 +197,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");
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

/**
Expand Down

0 comments on commit 4b97d8b

Please sign in to comment.