[remote/downloader] Wire credential helper to repository downloads

Progress on bazelbuild#16595
EngFlow · Nov 7, 2022 · 33299af · 33299af
1 parent 38eb886
commit 33299af
Show file tree

Hide file tree

Showing 5 changed files with 71 additions and 3 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/authandtls/GoogleAuthUtils.java b/src/main/java/com/google/devtools/build/lib/authandtls/GoogleAuthUtils.java
@@ -367,7 +367,6 @@ static Optional<Credentials> newCredentialsFromNetrc(
     }
   }
 
-  @VisibleForTesting
   public static CredentialHelperProvider newCredentialHelperProvider(
       CredentialHelperEnvironment environment,
       CommandLinePathFactory pathFactory,

diff --git a/src/main/java/com/google/devtools/build/lib/bazel/BUILD b/src/main/java/com/google/devtools/build/lib/bazel/BUILD
@@ -26,6 +26,8 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
         "//src/main/java/com/google/devtools/build/lib/analysis:blaze_directories",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/build_configuration",
+        "//src/main/java/com/google/devtools/build/lib/authandtls",
+        "//src/main/java/com/google/devtools/build/lib/authandtls/credentialhelper",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:common",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:inspection",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:inspection_impl",
@@ -54,6 +56,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi/repository",
         "//src/main/java/com/google/devtools/build/lib/util:abrupt_exit_exception",
         "//src/main/java/com/google/devtools/build/lib/util:detailed_exit_code",
+        "//src/main/java/com/google/devtools/build/lib/util:exit_code",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//src/main/java/com/google/devtools/common/options",

diff --git a/src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java b/src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java
@@ -15,6 +15,7 @@
 
 package com.google.devtools.build.lib.bazel;
 
+import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
 import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableList;
@@ -25,6 +26,12 @@
 import com.google.devtools.build.lib.analysis.ConfiguredRuleClassProvider;
 import com.google.devtools.build.lib.analysis.RuleDefinition;
 import com.google.devtools.build.lib.analysis.config.BuildConfigurationValue;
+import com.google.devtools.build.lib.authandtls.AuthAndTLSOptions;
+import com.google.devtools.build.lib.authandtls.GoogleAuthUtils;
+import com.google.devtools.build.lib.authandtls.StaticCredentials;
+import com.google.devtools.build.lib.authandtls.credentialhelper.CredentialHelperCredentials;
+import com.google.devtools.build.lib.authandtls.credentialhelper.CredentialHelperEnvironment;
+import com.google.devtools.build.lib.authandtls.credentialhelper.CredentialHelperProvider;
 import com.google.devtools.build.lib.bazel.bzlmod.BazelModuleInspectorFunction;
 import com.google.devtools.build.lib.bazel.bzlmod.BazelModuleInspectorValue.AugmentedModule.ResolutionReason;
 import com.google.devtools.build.lib.bazel.bzlmod.BazelModuleResolutionFunction;
@@ -78,6 +85,7 @@
 import com.google.devtools.build.lib.runtime.RepositoryRemoteExecutorFactory;
 import com.google.devtools.build.lib.runtime.ServerBuilder;
 import com.google.devtools.build.lib.runtime.WorkspaceBuilder;
+import com.google.devtools.build.lib.server.FailureDetails;
 import com.google.devtools.build.lib.server.FailureDetails.ExternalRepository;
 import com.google.devtools.build.lib.server.FailureDetails.ExternalRepository.Code;
 import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
@@ -90,6 +98,7 @@
 import com.google.devtools.build.lib.starlarkbuildapi.repository.RepositoryBootstrap;
 import com.google.devtools.build.lib.util.AbruptExitException;
 import com.google.devtools.build.lib.util.DetailedExitCode;
+import com.google.devtools.build.lib.util.ExitCode;
 import com.google.devtools.build.lib.vfs.FileSystem;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
@@ -363,6 +372,41 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
                 Code.BAD_DOWNLOADER_CONFIG));
       }
 
+      try {
+        AuthAndTLSOptions authAndTlsOptions = env.getOptions().getOptions(AuthAndTLSOptions.class);
+        var credentialHelperEnvironment =
+            CredentialHelperEnvironment.newBuilder()
+                .setEventReporter(env.getReporter())
+                .setWorkspacePath(env.getWorkspace())
+                .setClientEnvironment(env.getClientEnv())
+                .setHelperExecutionTimeout(authAndTlsOptions.credentialHelperTimeout)
+                .build();
+        CredentialHelperProvider credentialHelperProvider =
+            GoogleAuthUtils.newCredentialHelperProvider(
+                credentialHelperEnvironment,
+                env.getCommandLinePathFactory(),
+                authAndTlsOptions.credentialHelpers);
+
+        downloadManager.setCredentialFactory(headers -> {
+          Preconditions.checkNotNull(headers);
+
+          return new CredentialHelperCredentials(
+              credentialHelperProvider,
+              credentialHelperEnvironment,
+              Optional.of(new StaticCredentials(headers)),
+              authAndTlsOptions.credentialHelperCacheTimeout);
+        });
+      } catch (IOException e) {
+        env.getReporter().handle(Event.error(e.getMessage()));
+        env.getBlazeModuleEnvironment()
+            .exit(
+                new AbruptExitException(
+                    detailedExitCode(
+                        "Error initializing credential helper",
+                        Code.CREDENTIALS_INIT_FAILURE)));
+        return;
+      }
+
       if (repoOptions.experimentalDistdir != null) {
         downloadManager.setDistdir(
             repoOptions.experimentalDistdir.stream()

diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/downloader/DownloadManager.java b/src/main/java/com/google/devtools/build/lib/bazel/repository/downloader/DownloadManager.java
@@ -20,6 +20,7 @@
 import com.google.auth.Credentials;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Optional;
+import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -61,6 +62,7 @@ public class DownloadManager {
   private int retries = 0;
   private boolean urlsAsDefaultCanonicalId;
   @Nullable private Credentials netrcCreds;
+  private CredentialFactory credentialFactory = new DefaultCredentialFactory();
 
   public DownloadManager(RepositoryCache repositoryCache, Downloader downloader) {
     this.repositoryCache = repositoryCache;
@@ -92,6 +94,12 @@ public void setNetrcCreds(Credentials netrcCreds) {
     this.netrcCreds = netrcCreds;
   }
 
+  public void setCredentialFactory(CredentialFactory credentialFactory) {
+    Preconditions.checkNotNull(credentialFactory);
+
+    this.credentialFactory = credentialFactory;
+  }
+
   /**
    * Downloads file to disk and returns path.
    *
@@ -257,7 +265,7 @@ public Path download(
       try {
         downloader.download(
             rewrittenUrls,
-            new StaticCredentials(rewrittenAuthHeaders),
+            credentialFactory.create(rewrittenAuthHeaders),
             checksum,
             canonicalId,
             destination,
@@ -338,7 +346,7 @@ public byte[] downloadAndReadOneUrl(
     for (int attempt = 0; attempt <= retries; ++attempt) {
       try {
         return httpDownloader.downloadAndReadOneUrl(
-            rewrittenUrls.get(0), new StaticCredentials(authHeaders), eventHandler, clientEnv);
+            rewrittenUrls.get(0), credentialFactory.create(authHeaders), eventHandler, clientEnv);
       } catch (ContentLengthMismatchException e) {
         if (attempt == retries) {
           throw e;
@@ -427,4 +435,17 @@ public boolean isFinished() {
       return isFinished;
     }
   }
+
+  public interface CredentialFactory {
+    Credentials create(Map<URI, Map<String, List<String>>> authHeaders);
+  }
+
+  private static final class DefaultCredentialFactory implements CredentialFactory {
+    @Override
+    public Credentials create(Map<URI, Map<String, List<String>>> authHeaders) {
+      Preconditions.checkNotNull(authHeaders);
+
+      return new StaticCredentials(authHeaders);
+    }
+  }
 }
diff --git a/src/main/protobuf/failure_details.proto b/src/main/protobuf/failure_details.proto
@@ -248,6 +248,7 @@ message ExternalRepository {
     OVERRIDE_DISALLOWED_MANAGED_DIRECTORIES = 1 [(metadata) = { exit_code: 2 }];
     BAD_DOWNLOADER_CONFIG = 2 [(metadata) = { exit_code: 2 }];
     REPOSITORY_MAPPING_RESOLUTION_FAILED = 3 [(metadata) = { exit_code: 37 }];
+    CREDENTIALS_INIT_FAILURE = 4 [(metadata) = { exit_code: 2 }];
   }
   Code code = 1;
   // Additional data could include external repository names.