From 0e9ae8d360d3b45948fff504451130b72f335ed2 Mon Sep 17 00:00:00 2001
From: Brandon Duffany <brandon@buildbuddy.io>
Date: Mon, 13 Jun 2022 17:12:15 -0400
Subject: [PATCH] Fix bazel run argument handling (#2134)

---
 enterprise/server/cmd/ci_runner/main.go       | 25 ++++++++++--
 .../server/test/integration/ci_runner/BUILD   |  2 +-
 .../integration/ci_runner/ci_runner_test.go   | 40 +++++++++++++++++++
 3 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/enterprise/server/cmd/ci_runner/main.go b/enterprise/server/cmd/ci_runner/main.go
index dd02f0d3983..1792b96d872 100644
--- a/enterprise/server/cmd/ci_runner/main.go
+++ b/enterprise/server/cmd/ci_runner/main.go
@@ -760,7 +760,7 @@ func (ar *actionRunner) Run(ctx context.Context, ws *workspace) error {
 		// Transparently set the invocation ID from the one we computed ahead of
 		// time. The UI is expecting this invocation ID so that it can render a
 		// BuildBuddy invocation URL for each bazel_command that is executed.
-		args = append(args, fmt.Sprintf("--invocation_id=%s", iid))
+		args = appendBazelSubcommandArgs(args, fmt.Sprintf("--invocation_id=%s", iid))
 
 		// Instead of actually running the target, have Bazel write out a run script using the --script_path flag and
 		// extract run options (i.e. args, runfile information) from the generated run script.
@@ -772,7 +772,7 @@ func (ar *actionRunner) Run(ctx context.Context, ws *workspace) error {
 			}
 			defer os.RemoveAll(tmpDir)
 			runScript = filepath.Join(tmpDir, "run.sh")
-			args = append(args, "--script_path="+runScript)
+			args = appendBazelSubcommandArgs(args, "--script_path="+runScript)
 		}
 
 		runErr := runCommand(ctx, *bazelCommand, args, nil /*=env*/, ar.reporter)
@@ -1100,11 +1100,30 @@ func bazelArgs(cmd string) ([]string, error) {
 		if err != nil {
 			return nil, err
 		}
-		tokens = append(tokens, extras...)
+		tokens = appendBazelSubcommandArgs(tokens, extras...)
 	}
 	return append(startupFlags, tokens...), nil
 }
 
+// appendBazelSubcommandArgs appends bazel arguments to a bazel command,
+// *before* the arg separator ("--") if it exists, so that the arguments apply
+// to the bazel subcommand ("build", "run", etc.) and not the binary being run
+// (in the "bazel run" case).
+func appendBazelSubcommandArgs(args []string, argsToAppend ...string) []string {
+	splitIndex := len(args)
+	for i, arg := range args {
+		if arg == "--" {
+			splitIndex = i
+			break
+		}
+	}
+	out := make([]string, 0, len(args)+len(argsToAppend))
+	out = append(out, args[:splitIndex]...)
+	out = append(out, argsToAppend...)
+	out = append(out, args[splitIndex:]...)
+	return out
+}
+
 func ensureHomeDir() error {
 	// HOME dir is "/" when there is no user home dir created, which can happen
 	// when running locally (due to docker_inherit_user_ids). Treat this the same
diff --git a/enterprise/server/test/integration/ci_runner/BUILD b/enterprise/server/test/integration/ci_runner/BUILD
index 14bcea9d700..417a58a73d0 100644
--- a/enterprise/server/test/integration/ci_runner/BUILD
+++ b/enterprise/server/test/integration/ci_runner/BUILD
@@ -11,7 +11,7 @@ go_test(
     exec_properties = {
         "container-image": "docker://gcr.io/flame-public/buildbuddy-ci-runner:v2.2.8",
     },
-    shard_count = 8,
+    shard_count = 9,
     visibility = [
         "//enterprise:__subpackages__",
         "@buildbuddy_internal//enterprise:__subpackages__",
diff --git a/enterprise/server/test/integration/ci_runner/ci_runner_test.go b/enterprise/server/test/integration/ci_runner/ci_runner_test.go
index 9b072e521e1..1c0c051f705 100644
--- a/enterprise/server/test/integration/ci_runner/ci_runner_test.go
+++ b/enterprise/server/test/integration/ci_runner/ci_runner_test.go
@@ -70,6 +70,21 @@ actions:
 `,
 	}
 
+	workspaceContentsWithRunScript = map[string]string{
+		"WORKSPACE":     `workspace(name = "test")`,
+		"BUILD":         `sh_test(name = "print_args", srcs = ["print_args.sh"])`,
+		"print_args.sh": "echo 'args: {{' $@ '}}'",
+		"buildbuddy.yaml": `
+actions:
+  - name: "Print args"
+    triggers:
+      pull_request: { branches: [ master ] }
+      push: { branches: [ master ] }
+    bazel_commands:
+      - run //:print_args -- "Hello world"
+`,
+	}
+
 	invocationIDPattern = regexp.MustCompile(`Invocation URL:\s+.*?/invocation/([a-f0-9-]+)`)
 )
 
@@ -471,6 +486,31 @@ func TestCIRunner_PullRequest_FailedSync_CanRecoverAndRunCommand(t *testing.T) {
 	run()
 }
 
+func TestRunAction_RespectsArgs(t *testing.T) {
+	wsPath := testfs.MakeTempDir(t)
+	repoPath, headCommitSHA := makeGitRepo(t, workspaceContentsWithRunScript)
+
+	runnerFlags := []string{
+		"--workflow_id=test-workflow",
+		"--action_name=Print args",
+		"--trigger_event=push",
+		"--pushed_repo_url=file://" + repoPath,
+		"--pushed_branch=master",
+		"--commit_sha=" + headCommitSHA,
+		"--target_repo_url=file://" + repoPath,
+		"--target_branch=master",
+	}
+	// Start the app so the runner can use it as the BES backend.
+	app := buildbuddy.Run(t)
+	runnerFlags = append(runnerFlags, app.BESBazelFlags()...)
+
+	result := invokeRunner(t, runnerFlags, []string{}, wsPath)
+
+	checkRunnerResult(t, result)
+
+	assert.Contains(t, result.Output, "args: {{ Hello world }}")
+}
+
 func TestHostedBazel_ApplyingAndDiscardingPatches(t *testing.T) {
 	wsPath := testfs.MakeTempDir(t)