Add a test for RBE with executors talking to a Cache Proxy (#11342)

Related issues:
buildbuddy-io/buildbuddy-internal#4539

Author: iain-macdonald

enterprise/server/test/integration/cache_proxy/BUILD

@@ -12,7 +12,9 @@ go_test(
         "//proto:encryption_go_proto",
         "//proto:remote_execution_go_proto",
         "//server/testutil/testkeys",
+        "//server/util/authutil",
         "//server/util/claims",
+        "//server/util/grpc_client",
         "//server/util/grpc_server",
         "//server/util/testing/flags",
         "@com_github_golang_jwt_jwt_v4//:jwt",

enterprise/server/test/integration/cache_proxy/cache_proxy_test.go

@@ -2,13 +2,17 @@ package cache_proxy_test
 
 import (
 	"context"
+	"fmt"
+	"runtime"
 	"strings"
 	"sync"
 	"testing"
 
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/test/integration/remote_execution/rbetest"
 	"github.com/buildbuddy-io/buildbuddy/server/testutil/testkeys"
+	"github.com/buildbuddy-io/buildbuddy/server/util/authutil"
 	"github.com/buildbuddy-io/buildbuddy/server/util/claims"
+	"github.com/buildbuddy-io/buildbuddy/server/util/grpc_client"
 	"github.com/buildbuddy-io/buildbuddy/server/util/grpc_server"
 	"github.com/buildbuddy-io/buildbuddy/server/util/testing/flags"
 	"github.com/golang-jwt/jwt/v4"
@@ -40,13 +44,14 @@ func TestES256Auth(t *testing.T) {
 	keyPair := testkeys.GenerateES256KeyPair(t)
 	flags.Set(t, "auth.jwt_es256_private_key", keyPair.PrivateKeyPEM)
 	flags.Set(t, "auth.remote.use_es256_jwts", true)
+	flags.Set(t, "auth.reparse_jwts", false)
 	require.NoError(t, claims.Init())
 
 	var mu sync.Mutex
 	var capturedJWT string
 	jwtInterceptor := func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
 		if md, ok := metadata.FromIncomingContext(ctx); ok {
-			if vals := md.Get("x-buildbuddy-jwt"); len(vals) > 0 {
+			if vals := md.Get(authutil.ContextTokenStringKey); len(vals) > 0 {
 				mu.Lock()
 				capturedJWT = vals[len(vals)-1]
 				mu.Unlock()
@@ -69,7 +74,7 @@ func TestES256Auth(t *testing.T) {
 		},
 	}
 
-	ctx := metadata.AppendToOutgoingContext(t.Context(), "x-buildbuddy-api-key", rbe.APIKey1)
+	ctx := metadata.AppendToOutgoingContext(t.Context(), authutil.APIKeyHeader, rbe.APIKey1)
 	resp, err := cas.FindMissingBlobs(ctx, &req)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(resp.MissingBlobDigests))
@@ -84,6 +89,42 @@ func TestES256Auth(t *testing.T) {
 	require.Equal(t, "ES256", token.Method.Alg())
 }
 
+func TestES256Auth_RemoteExecution(t *testing.T) {
+	// TODO(http://go/b/4539): enable this test
+	t.Skip()
+
+	keyPair := testkeys.GenerateES256KeyPair(t)
+	flags.Set(t, "auth.jwt_es256_private_key", keyPair.PrivateKeyPEM)
+	flags.Set(t, "auth.remote.use_es256_jwts", true)
+	flags.Set(t, "auth.reparse_jwts", false)
+	require.NoError(t, claims.Init())
+
+	rbe := rbetest.NewRBETestEnv(t)
+	rbe.AddBuildBuddyServer()
+	proxy := rbe.AddCacheProxy()
+	conn, err := grpc_client.DialSimple(
+		fmt.Sprintf("grpc://localhost:%d", proxy.Port))
+	require.NoError(t, err)
+	rbe.AddExecutorWithOptions(t, &rbetest.ExecutorOptions{
+		Name:      "executor",
+		CacheConn: conn,
+	})
+
+	cmd := rbe.Execute(&repb.Command{
+		Arguments: []string{"sh", "-c", "echo hello"},
+		Platform: &repb.Platform{
+			Properties: []*repb.Platform_Property{
+				{Name: "container-image", Value: "none"},
+				{Name: "OSFamily", Value: runtime.GOOS},
+				{Name: "Arch", Value: runtime.GOARCH},
+			},
+		},
+	}, &rbetest.ExecuteOpts{APIKey: rbe.APIKey1})
+	res := cmd.Wait()
+	require.Equal(t, 0, res.ExitCode)
+	require.Equal(t, "hello\n", res.Stdout)
+}
+
 func TestFindMissing_Encryption(t *testing.T) {
 	rbe := rbetest.NewRBETestEnv(t)
 	rbe.AddBuildBuddyServer()
@@ -97,7 +138,7 @@ func TestFindMissing_Encryption(t *testing.T) {
 		},
 	}
 
-	ctx := metadata.AppendToOutgoingContext(t.Context(), "x-buildbuddy-api-key", rbe.APIKey1)
+	ctx := metadata.AppendToOutgoingContext(t.Context(), authutil.APIKeyHeader, rbe.APIKey1)
 	resp, err := cas.FindMissingBlobs(ctx, &req)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(resp.MissingBlobDigests))

enterprise/server/test/integration/remote_execution/rbetest/rbetest.go

@@ -710,6 +710,9 @@ type ExecutorOptions struct {
 	APIKey string
 	// Optional Pool name for the executor
 	Pool string
+	// Optional connection for the executor to use for CAS/AC/ByteStream.
+	// Defaults to the app proxy connection if unset.
+	CacheConn grpc.ClientConnInterface
 	// Optional interceptor for command execution results.
 	RunInterceptor
 	priorityTaskSchedulerOptions priority_task_scheduler.Options
@@ -859,11 +862,16 @@ func (r *Env) addExecutor(t testing.TB, options *ExecutorOptions) *Executor {
 	clientConn := r.appProxyConn
 	env.SetSchedulerClient(scpb.NewSchedulerClient(clientConn))
 	env.SetRemoteExecutionClient(repb.NewExecutionClient(clientConn))
-	env.SetActionCacheClient(repb.NewActionCacheClient(clientConn))
-	env.SetContentAddressableStorageClient(repb.NewContentAddressableStorageClient(clientConn))
-	env.SetByteStreamClient(bspb.NewByteStreamClient(clientConn))
 	env.SetCapabilitiesClient(repb.NewCapabilitiesClient(clientConn))
 
+	cacheConn := grpc.ClientConnInterface(clientConn)
+	if options.CacheConn != nil {
+		cacheConn = options.CacheConn
+	}
+	env.SetActionCacheClient(repb.NewActionCacheClient(cacheConn))
+	env.SetContentAddressableStorageClient(repb.NewContentAddressableStorageClient(cacheConn))
+	env.SetByteStreamClient(bspb.NewByteStreamClient(cacheConn))
+
 	env.SetAuthenticator(r.testEnv.GetAuthenticator())
 	xl := xcode.NewXcodeLocator()
 	env.SetXcodeLocator(xl)
@@ -976,7 +984,7 @@ func (r *Env) waitForExecutorRegistration() {
 type CacheProxy struct {
 	t    testing.TB
 	env  *testenv.TestEnv
-	port int
+	Port int
 	conn *grpc_client.ClientConnPool
 }
 
@@ -1036,7 +1044,7 @@ func (r *Env) AddCacheProxy() *CacheProxy {
 	// Finally, create the client connection.
 	conn, err := grpc_client.DialSimple(fmt.Sprintf("grpc://localhost:%d", port))
 	require.NoError(r.t, err)
-	return &CacheProxy{t: r.t, env: proxyEnv, conn: conn}
+	return &CacheProxy{t: r.t, env: proxyEnv, Port: port, conn: conn}
 }
 
 func (r *Env) DownloadOutputsToNewTempDir(res *CommandResult) string {

server/testutil/testauth/testauth.go

@@ -219,7 +219,13 @@ func RequestContext(userID string, groupID string) *ctxpb.RequestContext {
 
 // WithAuthenticatedUserInfo sets the authenticated user to the given user.
 func WithAuthenticatedUserInfo(ctx context.Context, userInfo interfaces.UserInfo) context.Context {
-	jwt, err := claims.AssembleJWT(userInfo.(*claims.Claims), jwt.SigningMethodHS256)
+	// Copy claims to avoid mutating shared state when AssembleJWT sets standard
+	// claims, then set them in context to avoid reparsing.
+	newClaims := *userInfo.(*claims.Claims)
+	c := &newClaims
+	ctx = claims.AuthContext(ctx, c)
+
+	jwt, err := claims.AssembleJWT(c, jwt.SigningMethodHS256)
 	if err != nil {
 		log.Errorf("Failed to mint JWT from UserInfo: %s", err)
 		return ctx