Slap the parsed jwt on the context once instead of re-parsing ad naus…

…eum (#8032)
buildbuddy-io · Dec 9, 2024 · 31a1518 · 31a1518
1 parent 95ed772
commit 31a1518
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git a/enterprise/server/remote_execution/executor/BUILD b/enterprise/server/remote_execution/executor/BUILD
@@ -20,8 +20,10 @@ go_library(
         "//server/metrics",
         "//server/remote_cache/cachetools",
         "//server/remote_cache/digest",
+        "//server/rpc/interceptors",
         "//server/util/alert",
         "//server/util/background",
+        "//server/util/bazel_request",
         "//server/util/canary",
         "//server/util/disk",
         "//server/util/log",

diff --git a/enterprise/server/remote_execution/executor/executor.go b/enterprise/server/remote_execution/executor/executor.go
@@ -19,8 +19,10 @@ import (
 	"github.com/buildbuddy-io/buildbuddy/server/metrics"
 	"github.com/buildbuddy-io/buildbuddy/server/remote_cache/cachetools"
 	"github.com/buildbuddy-io/buildbuddy/server/remote_cache/digest"
+	"github.com/buildbuddy-io/buildbuddy/server/rpc/interceptors"
 	"github.com/buildbuddy-io/buildbuddy/server/util/alert"
 	"github.com/buildbuddy-io/buildbuddy/server/util/background"
+	"github.com/buildbuddy-io/buildbuddy/server/util/bazel_request"
 	"github.com/buildbuddy-io/buildbuddy/server/util/canary"
 	"github.com/buildbuddy-io/buildbuddy/server/util/disk"
 	"github.com/buildbuddy-io/buildbuddy/server/util/log"
@@ -179,6 +181,9 @@ func (s *Executor) ExecuteTaskAndStreamResults(ctx context.Context, st *repb.Sch
 	ctx, span := tracing.StartSpan(ctx)
 	defer span.End()
 
+	ctx = interceptors.AddAuthToContext(s.env, ctx)
+	ctx = bazel_request.ParseRequestMetadataOnce(ctx)
+
 	metrics.RemoteExecutionTasksStartedCount.Inc()
 
 	actionMetrics := &ActionMetrics{}

diff --git a/server/rpc/interceptors/interceptors.go b/server/rpc/interceptors/interceptors.go
@@ -91,7 +91,7 @@ func contextReplacingUnaryClientInterceptor(ctxFn func(ctx context.Context) cont
 	}
 }
 
-func addAuthToContext(env environment.Env, ctx context.Context) context.Context {
+func AddAuthToContext(env environment.Env, ctx context.Context) context.Context {
 	ctx = env.GetAuthenticator().AuthenticatedGRPCContext(ctx)
 	if c, err := claims.ClaimsFromContext(ctx); err == nil {
 		ctx = log.EnrichContext(ctx, "group_id", c.GetGroupID())
@@ -163,7 +163,7 @@ func setHeadersFromContext(ctx context.Context) context.Context {
 // middleware to the request.
 func authStreamServerInterceptor(env environment.Env) grpc.StreamServerInterceptor {
 	ctxFn := func(ctx context.Context) context.Context {
-		return addAuthToContext(env, ctx)
+		return AddAuthToContext(env, ctx)
 	}
 	return contextReplacingStreamServerInterceptor(ctxFn)
 }
@@ -172,7 +172,7 @@ func authStreamServerInterceptor(env environment.Env) grpc.StreamServerIntercept
 // middleware to the request.
 func authUnaryServerInterceptor(env environment.Env) grpc.UnaryServerInterceptor {
 	ctxFn := func(ctx context.Context) context.Context {
-		return addAuthToContext(env, ctx)
+		return AddAuthToContext(env, ctx)
 	}
 	return contextReplacingUnaryServerInterceptor(ctxFn)
 }