[RB] Support passing a different API key for uploading build events

Author: maggie-lou

app/invocation/invocation_action_card.tsx

@@ -657,7 +657,10 @@ export default class InvocationActionCardComponent extends React.Component<Props
   ) {
     const snapshotKey = this.getSnapshotKeyForSnapshotID(vmMetadata);
     const snapshotKeyJSON = JSON.stringify(snapshotKey);
-    const cmd = `bb remote --run_from_snapshot='${snapshotKeyJSON}' --runner_exec_properties=debug-executor-id=${executionMetadata.executorId} --script='echo "My custom bash command!"'`;
+    const impersonationAPIKey = this.state.user?.isImpersonating
+      ? `--remote_run_header=x-buildbuddy-platform.env-overrides=BUILDBUDDY_BES_API_KEY=[SET BB ORG API KEY]>`
+      : "";
+    const cmd = `bb remote ${impersonationAPIKey} --run_from_snapshot='${snapshotKeyJSON}' --runner_exec_properties=debug-executor-id=${executionMetadata.executorId} --script='echo "My custom bash command!"'`;
     copyToClipboard(cmd);
     alert_service.success("Command copied to clipboard");
     this.setState({ showSnapshotMenu: false });

enterprise/server/cmd/ci_runner/main.go

@@ -680,9 +680,13 @@ func run() error {
 		defer cancel()
 	}
 
+	besAPIKey := os.Getenv("BUILDBUDDY_BES_API_KEY")
+	if besAPIKey == "" {
+		besAPIKey = ws.buildbuddyAPIKey
+	}
 	// Use a context without a timeout for the build event reporter, so that even
 	// if the `timeout` is reached, any events will finish getting published
-	buildEventReporter, err := newBuildEventReporter(contextWithoutTimeout, *besBackend, ws.buildbuddyAPIKey, *invocationID, *workflowID != "" /*=isWorkflow*/)
+	buildEventReporter, err := newBuildEventReporter(contextWithoutTimeout, *besBackend, besAPIKey, *invocationID, *workflowID != "" /*=isWorkflow*/)
 	if err != nil {
 		return err
 	}

enterprise/server/hostedrunner/BUILD

@@ -11,6 +11,8 @@ go_library(
         "//enterprise/server/remote_execution/platform",
         "//enterprise/server/util/ci_runner_util",
         "//enterprise/server/workflow/config",
+        "//proto:auditlog_go_proto",
+        "//proto:firecracker_go_proto",
         "//proto:remote_execution_go_proto",
         "//proto:runner_go_proto",
         "//server/endpoint_urls/build_buddy_url",
@@ -33,6 +35,7 @@ go_library(
         "@org_golang_google_genproto//googleapis/longrunning",
         "@org_golang_google_grpc//metadata",
         "@org_golang_google_grpc//status",
+        "@org_golang_google_protobuf//encoding/protojson",
         "@org_golang_google_protobuf//types/known/durationpb",
     ],
 )

enterprise/server/hostedrunner/hostedrunner.go

@@ -12,6 +12,7 @@ import (
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/remote_execution/platform"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/util/ci_runner_util"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/workflow/config"
+	"github.com/buildbuddy-io/buildbuddy/proto/auditlog"
 	"github.com/buildbuddy-io/buildbuddy/server/endpoint_urls/build_buddy_url"
 	"github.com/buildbuddy-io/buildbuddy/server/endpoint_urls/cache_api_url"
 	"github.com/buildbuddy-io/buildbuddy/server/endpoint_urls/events_api_url"
@@ -30,9 +31,11 @@ import (
 	"github.com/google/uuid"
 	"google.golang.org/genproto/googleapis/longrunning"
 	"google.golang.org/grpc/metadata"
+	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/types/known/durationpb"
 	"gopkg.in/yaml.v2"
 
+	fcpb "github.com/buildbuddy-io/buildbuddy/proto/firecracker"
 	repb "github.com/buildbuddy-io/buildbuddy/proto/remote_execution"
 	rnpb "github.com/buildbuddy-io/buildbuddy/proto/runner"
 	gstatus "google.golang.org/grpc/status"
@@ -420,6 +423,7 @@ func (r *runnerService) Run(ctx context.Context, req *rnpb.RunRequest) (*rnpb.Ru
 		return nil, status.WrapError(err, "get credentials")
 	}
 
+	var hasBESOverride bool
 	for _, h := range req.GetRemoteHeaders() {
 		parts := strings.SplitN(h, "=", 2)
 		if len(parts) != 2 {
@@ -432,6 +436,10 @@ func (r *runnerService) Run(ctx context.Context, req *rnpb.RunRequest) (*rnpb.Ru
 		// to credential-related env overrides that were set above.
 		if headerKey == platform.OverrideHeaderPrefix+platform.EnvOverridesPropertyName {
 			envOverrides = append(envOverrides, headerVal)
+
+			if strings.HasPrefix(headerVal, "BUILDBUDDY_BES_API_KEY=") {
+				hasBESOverride = true
+			}
 			continue
 		}
 
@@ -441,6 +449,16 @@ func (r *runnerService) Run(ctx context.Context, req *rnpb.RunRequest) (*rnpb.Ru
 	execCtx = platform.WithRemoteHeaderOverride(
 		execCtx, platform.EnvOverridesPropertyName, strings.Join(envOverrides, ","))
 
+	if hasBESOverride && r.env.GetAuditLogger() != nil {
+		snapshotKeyStr := getExecProperty(req.GetExecProperties(), "snapshot-key-override")
+		snapshotKey := &fcpb.SnapshotKey{}
+		if err := protojson.Unmarshal([]byte(snapshotKeyStr), snapshotKey); err != nil {
+			return nil, status.WrapError(err, "unmarshal SnapshotKey")
+		}
+		// TODO: Make sure this is logging for the correct group
+		r.env.GetAuditLogger().LogForGroup(ctx, req.GetRequestContext().GetGroupId(), auditlog.Action_IMPERSONATE_REMOTE_RUN, snapshotKey)
+	}
+
 	executionClient := r.env.GetRemoteExecutionClient()
 	if executionClient == nil {
 		return nil, status.UnimplementedError("Missing remote execution client.")

proto/BUILD

@@ -56,6 +56,7 @@ proto_library(
         ":api_key_proto",
         ":context_proto",
         ":encryption_proto",
+        ":firecracker_proto",
         ":github_proto",
         ":group_proto",
         ":invocation_proto",
@@ -926,6 +927,7 @@ go_proto_library(
         ":api_key_go_proto",
         ":context_go_proto",
         ":encryption_go_proto",
+        ":firecracker_go_proto",
         ":github_go_proto",
         ":group_go_proto",
         ":invocation_go_proto",
@@ -2070,6 +2072,7 @@ ts_proto_library(
         ":api_key_ts_proto",
         ":context_ts_proto",
         ":encryption_ts_proto",
+        ":firecracker_ts_proto",
         ":github_ts_proto",
         ":group_ts_proto",
         ":invocation_ts_proto",

proto/auditlog.proto

@@ -5,6 +5,7 @@ package auditlog;
 import "proto/api_key.proto";
 import "proto/context.proto";
 import "proto/encryption.proto";
+import "proto/firecracker.proto";
 import "proto/github.proto";
 import "proto/grp.proto";
 import "proto/invocation.proto";
@@ -61,6 +62,7 @@ enum Action {
   CREATE_IMPERSONATION_API_KEY = 12;
   UPDATE_IP_RULES_CONFIG = 13;
   INVALIDATE_VM_SNAPSHOT = 14;
+  IMPERSONATE_REMOTE_RUN = 15;
 }
 
 message ResourceID {
@@ -101,6 +103,7 @@ message Entry {
     iprules.DeleteRuleRequest delete_ip_rule = 17;
     iprules.SetRulesConfigRequest set_rules_config = 18;
     workflow.InvalidateSnapshotRequest invalidate_snapshot = 19;
+    firecracker.SnapshotKey impersonate_remote_run = 20;
   }
   message Request {
     APIRequest api_request = 1;