From 6ff52dc457898c1e884d8be7a05a3ffc3fb9845d Mon Sep 17 00:00:00 2001
From: Fabian Meumertzheim <fabian@meumertzhe.im>
Date: Thu, 7 Nov 2024 02:58:49 -0800
Subject: [PATCH] Let repo rules remove env vars for subprocesses

RELNOTES: `repository_ctx.execute` can now remove an environment variable when executing a process by associating it with the value `None` in the `environment` argument.

Closes #24221.

PiperOrigin-RevId: 694048622
Change-Id: Ia779e9133411745c9cfe93845c13839d05731c7c
---
 .../starlark/StarlarkBaseExternalContext.java | 33 +++++++++++++---
 .../starlark/StarlarkExecutionResult.java     |  8 ++++
 src/test/shell/bazel/BUILD                    |  1 +
 .../shell/bazel/starlark_repository_test.sh   | 38 +++++++++++++++++++
 4 files changed, 75 insertions(+), 5 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkBaseExternalContext.java b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkBaseExternalContext.java
index 8902fe24254db4..e34f582b977402 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkBaseExternalContext.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkBaseExternalContext.java
@@ -84,9 +84,12 @@
 import java.util.Base64;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -1835,7 +1838,11 @@ private static String profileArgsDesc(String method, List<String> args) {
             name = "environment",
             defaultValue = "{}",
             named = true,
-            doc = "Force some environment variables to be set to be passed to the process."),
+            doc =
+                """
+                Force some environment variables to be set to be passed to the process. The value \
+                can be <code>None</code> to remove the environment variable.
+                """),
         @Param(
             name = "quiet",
             defaultValue = "True",
@@ -1855,7 +1862,7 @@ private static String profileArgsDesc(String method, List<String> args) {
   public StarlarkExecutionResult execute(
       Sequence<?> arguments, // <String> or <StarlarkPath> or <Label> expected
       StarlarkInt timeoutI,
-      Dict<?, ?> uncheckedEnvironment, // <String, String> expected
+      Dict<?, ?> uncheckedEnvironment, // <String, Object> expected
       boolean quiet,
       String overrideWorkingDirectory,
       StarlarkThread thread)
@@ -1863,10 +1870,25 @@ public StarlarkExecutionResult execute(
     validateExecuteArguments(arguments);
     int timeout = Starlark.toInt(timeoutI, "timeout");
 
-    Map<String, String> forceEnvVariables =
-        Dict.cast(uncheckedEnvironment, String.class, String.class, "environment");
+    Map<String, Object> forceEnvVariablesRaw =
+        Dict.cast(uncheckedEnvironment, String.class, Object.class, "environment");
+    Map<String, String> forceEnvVariables = new LinkedHashMap<>();
+    Set<String> removeEnvVariables = new LinkedHashSet<>();
+    for (Map.Entry<String, Object> entry : forceEnvVariablesRaw.entrySet()) {
+      String key = entry.getKey();
+      Object value = entry.getValue();
+      if (value == Starlark.NONE) {
+        removeEnvVariables.add(key);
+      } else if (value instanceof String s) {
+        forceEnvVariables.put(key, s);
+      } else {
+        throw Starlark.errorf("environment values must be strings or None, got %s", value);
+      }
+    }
 
     if (canExecuteRemote()) {
+      // Remote execution only sees the explicitly set environment variables, so removing env vars
+      // isn't necessary.
       return executeRemote(arguments, timeout, forceEnvVariables, quiet, overrideWorkingDirectory);
     }
 
@@ -1886,7 +1908,7 @@ public StarlarkExecutionResult execute(
         WorkspaceRuleEvent.newExecuteEvent(
             args,
             timeout,
-            envVariables,
+            Maps.filterKeys(envVariables, k -> !removeEnvVariables.contains(k)),
             forceEnvVariables,
             workingDirectory.getPathString(),
             quiet,
@@ -1920,6 +1942,7 @@ public StarlarkExecutionResult execute(
           .addArguments(args)
           .setDirectory(workingDirectoryPath.getPathFile())
           .addEnvironmentVariables(forceEnvVariables)
+          .removeEnvironmentVariables(removeEnvVariables)
           .setTimeout(timeoutMillis)
           .setQuiet(quiet)
           .execute();
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkExecutionResult.java b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkExecutionResult.java
index eaa7d1d586b49e..7bd40f94599f20 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkExecutionResult.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkExecutionResult.java
@@ -35,6 +35,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.annot.StarlarkMethod;
 import net.starlark.java.eval.EvalException;
@@ -158,6 +159,13 @@ Builder addEnvironmentVariables(Map<String, String> variables) {
       return this;
     }
 
+    /** Ensure that an environment variable is not passed to the process. */
+    @CanIgnoreReturnValue
+    Builder removeEnvironmentVariables(Set<String> removeEnvVariables) {
+      removeEnvVariables.forEach(envBuilder::remove);
+      return this;
+    }
+
     /** Sets the timeout, in milliseconds, after which the executed command will be terminated. */
     @CanIgnoreReturnValue
     Builder setTimeout(long timeout) {
diff --git a/src/test/shell/bazel/BUILD b/src/test/shell/bazel/BUILD
index 4295ef2c8bf9df..cc5ebafe945635 100644
--- a/src/test/shell/bazel/BUILD
+++ b/src/test/shell/bazel/BUILD
@@ -874,6 +874,7 @@ sh_test(
     data = [
         ":test-deps",
         "@bazel_tools//tools/bash/runfiles",
+        "@local_jdk//:jdk",
     ],
     shard_count = 10,
     tags = [
diff --git a/src/test/shell/bazel/starlark_repository_test.sh b/src/test/shell/bazel/starlark_repository_test.sh
index 6ba32d122c09b3..aad44faedd0e66 100755
--- a/src/test/shell/bazel/starlark_repository_test.sh
+++ b/src/test/shell/bazel/starlark_repository_test.sh
@@ -3448,4 +3448,42 @@ EOF
     expect_log "my_second_repo: @my_first_repo//:foo @@+_repo_rules+my_first_repo//:foo"
 }
 
+function test_execute_environment_remove_vars() {
+  cat >> $(setup_module_dot_bazel)  <<'EOF'
+my_repo = use_repo_rule("//:repo.bzl", "my_repo")
+my_repo(name="repo")
+EOF
+  touch BUILD
+  cat > repo.bzl <<'EOF'
+def _impl(ctx):
+  st = ctx.execute(
+    ["env"],
+    environment = {
+      "CLIENT_ENV_REMOVED": None,
+      "REPO_ENV_REMOVED": None,
+    },
+  )
+  if st.return_code:
+    fail("Command did not succeed")
+  vars = {line.partition("=")[0]: line.partition("=")[-1] for line in st.stdout.split("\n")}
+  if "CLIENT_ENV_REMOVED" in vars:
+    fail("CLIENT_ENV_REMOVED should not be in the environment")
+  if "REPO_ENV_REMOVED" in vars:
+    fail("REPO_ENV_REMOVED should not be in the environment")
+  if vars.get("CLIENT_ENV_PRESENT") != "value1":
+    fail("CLIENT_ENV_PRESENT has wrong value: " + vars.get("CLIENT_ENV_PRESENT"))
+  if vars.get("REPO_ENV_PRESENT") != "value3":
+    fail("REPO_ENV_PRESENT has wrong value: " + vars.get("REPO_ENV_PRESENT"))
+
+  ctx.file("BUILD", "exports_files(['data.txt'])")
+
+my_repo = repository_rule(_impl)
+EOF
+
+  CLIENT_ENV_PRESENT=value1 CLIENT_ENV_REMOVED=value2 \
+   bazel build \
+    --repo_env=REPO_ENV_PRESENT=value3 --repo_env=REPO_ENV_REMOVED=value4 \
+    @repo//... &> $TEST_log || fail "expected Bazel to succeed"
+}
+
 run_suite "local repository tests"