Store yanked but allowed versions

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelLockFileModule.java

@@ -117,9 +117,10 @@ public void afterCommand() {
     // lockfile that are still up-to-date and adding the newly resolved extension results.
     BazelLockFileValue newLockfile =
         BazelLockFileValue.builder()
-            .setModuleExtensions(combinedExtensionInfos)
             .setRegistryFileHashes(
                 ImmutableSortedMap.copyOf(moduleResolutionValue.getRegistryFileHashes()))
+            .setYankedButAllowedModules(moduleResolutionValue.getYankedButAllowedModules())
+            .setModuleExtensions(combinedExtensionInfos)
             .build();
 
     // Write the new value to the file, but only if needed. This is not just a performance

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelLockFileValue.java

@@ -17,6 +17,7 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
 import com.google.devtools.build.lib.events.ExtendedEventHandler.Postable;
 import com.google.devtools.build.lib.skyframe.SkyFunctions;
@@ -43,6 +44,7 @@ static Builder builder() {
     return new AutoValue_BazelLockFileValue.Builder()
         .setLockFileVersion(LOCK_FILE_VERSION)
         .setRegistryFileHashes(ImmutableMap.of())
+        .setYankedButAllowedModules(ImmutableSet.of())
         .setModuleExtensions(ImmutableMap.of());
   }
 
@@ -52,6 +54,9 @@ static Builder builder() {
   /** Hashes of files retrieved from registries. */
   public abstract ImmutableMap<String, Optional<Checksum>> getRegistryFileHashes();
 
+  /** Module versions that are known to be yanked but were explicitly allowed by the user. */
+  public abstract ImmutableSet<ModuleKey> getYankedButAllowedModules();
+
   /** Mapping the extension id to the module extension data */
   public abstract ImmutableMap<
           ModuleExtensionId, ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
@@ -66,6 +71,8 @@ public abstract static class Builder {
 
     public abstract Builder setRegistryFileHashes(ImmutableMap<String, Optional<Checksum>> value);
 
+    public abstract Builder setYankedButAllowedModules(ImmutableSet<ModuleKey> value);
+
     public abstract Builder setModuleExtensions(
         ImmutableMap<
                 ModuleExtensionId,

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelModuleResolutionFunction.java

@@ -68,7 +68,8 @@ public class BazelModuleResolutionFunction implements SkyFunction {
 
   private record Result(
       Selection.Result selectionResult,
-      ImmutableMap<String, Optional<Checksum>> registryFileHashes) {}
+      ImmutableMap<String, Optional<Checksum>> registryFileHashes,
+      ImmutableSet<ModuleKey> yankedButAllowedModules) {}
 
   private static class ModuleResolutionComputeState implements Environment.SkyKeyComputeState {
     Result discoverAndSelectResult;
@@ -142,7 +143,8 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     return BazelModuleResolutionValue.create(
         finalDepGraph,
         state.discoverAndSelectResult.selectionResult.getUnprunedDepGraph(),
-        ImmutableMap.copyOf(registryFileHashes));
+        ImmutableMap.copyOf(registryFileHashes),
+        state.discoverAndSelectResult.yankedButAllowedModules);
   }
 
   @Nullable
@@ -206,12 +208,15 @@ private static Result discoverAndSelect(
           env.getListener());
     }
 
+    ImmutableSet<ModuleKey> yankedButAllowedModules;
     try (SilentCloseable c =
         Profiler.instance().profile(ProfilerTask.BZLMOD, "check no yanked versions")) {
-      checkNoYankedVersions(resolvedDepGraph, yankedVersionValues, allowedYankedVersions);
+      yankedButAllowedModules =
+          checkNoYankedVersions(resolvedDepGraph, yankedVersionValues, allowedYankedVersions);
     }
 
-    return new Result(selectionResult, discoveryResult.registryFileHashes());
+    return new Result(
+        selectionResult, discoveryResult.registryFileHashes(), yankedButAllowedModules);
   }
 
   private static void verifyAllOverridesAreOnExistentModules(
@@ -305,36 +310,53 @@ public static void checkBazelCompatibility(
     }
   }
 
-  private static void checkNoYankedVersions(
+  /**
+   * Fail if any selected module is yanked and not explicitly allowed.
+   *
+   * @return the set of yanked but explicitly allowed modules
+   */
+  private static ImmutableSet<ModuleKey> checkNoYankedVersions(
       ImmutableMap<ModuleKey, InterimModule> depGraph,
       ImmutableMap<YankedVersionsValue.Key, YankedVersionsValue> yankedVersionValues,
       Optional<ImmutableSet<ModuleKey>> allowedYankedVersions)
       throws BazelModuleResolutionFunctionException {
+    ImmutableSet.Builder<ModuleKey> yankedButAllowedModules = ImmutableSet.builder();
     for (InterimModule m : depGraph.values()) {
       if (m.getRegistry() == null) {
         // Non-registry modules do not have yanked versions.
         continue;
       }
-      Optional<String> yankedInfo =
-          YankedVersionsUtil.getYankedInfo(
-              m.getKey(),
-              yankedVersionValues.get(
-                  YankedVersionsValue.Key.create(m.getKey(), m.getRegistry().getUrl())),
-              allowedYankedVersions);
-      if (yankedInfo.isPresent()) {
-        throw new BazelModuleResolutionFunctionException(
-            ExternalDepsException.withMessage(
-                Code.VERSION_RESOLUTION_ERROR,
-                "Yanked version detected in your resolved dependency graph: %s, for the reason: "
-                    + "%s.\nYanked versions may contain serious vulnerabilities and should not be "
-                    + "used. To fix this, use a bazel_dep on a newer version of this module. To "
-                    + "continue using this version, allow it using the --allow_yanked_versions "
-                    + "flag or the BZLMOD_ALLOW_YANKED_VERSIONS env variable.",
-                m.getKey(),
-                yankedInfo.get()),
-            Transience.PERSISTENT);
+      ModuleKey key = m.getKey();
+      YankedVersionsValue yankedVersionsValue =
+          yankedVersionValues.get(
+              YankedVersionsValue.Key.create(m.getKey(), m.getRegistry().getUrl()));
+      if (yankedVersionsValue.yankedVersions().isEmpty()) {
+        // No yanked version information available or no need to check it.
+        continue;
+      }
+      String yankedInfo = yankedVersionsValue.yankedVersions().get().get(key.getVersion());
+      if (yankedInfo == null) {
+        // The version is not yanked.
+        continue;
+      }
+      if (allowedYankedVersions.isEmpty() || allowedYankedVersions.get().contains(key)) {
+        // The version is yanked but explicitly allowed.
+        yankedButAllowedModules.add(key);
+        continue;
       }
+      throw new BazelModuleResolutionFunctionException(
+          ExternalDepsException.withMessage(
+              Code.VERSION_RESOLUTION_ERROR,
+              "Yanked version detected in your resolved dependency graph: %s, for the reason: "
+                  + "%s.\nYanked versions may contain serious vulnerabilities and should not be "
+                  + "used. To fix this, use a bazel_dep on a newer version of this module. To "
+                  + "continue using this version, allow it using the --allow_yanked_versions "
+                  + "flag or the BZLMOD_ALLOW_YANKED_VERSIONS env variable.",
+              m.getKey(),
+              yankedInfo),
+          Transience.PERSISTENT);
     }
+    return yankedButAllowedModules.build();
   }
 
   private static ImmutableMap<ModuleKey, Module> computeFinalDepGraph(

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelModuleResolutionValue.java

@@ -17,6 +17,7 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
 import com.google.devtools.build.lib.skyframe.SkyFunctions;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationConstant;
@@ -50,11 +51,15 @@ abstract class BazelModuleResolutionValue implements SkyValue {
    */
   abstract ImmutableMap<String, Optional<Checksum>> getRegistryFileHashes();
 
+  /** Module versions that are known to be yanked but were explicitly allowed by the user. */
+  abstract ImmutableSet<ModuleKey> getYankedButAllowedModules();
+
   static BazelModuleResolutionValue create(
       ImmutableMap<ModuleKey, Module> resolvedDepGraph,
       ImmutableMap<ModuleKey, InterimModule> unprunedDepGraph,
-      ImmutableMap<String, Optional<Checksum>> registryFileHashes) {
+      ImmutableMap<String, Optional<Checksum>> registryFileHashes,
+      ImmutableSet<ModuleKey> yankedButAllowedModules) {
     return new AutoValue_BazelModuleResolutionValue(
-        resolvedDepGraph, unprunedDepGraph, registryFileHashes);
+        resolvedDepGraph, unprunedDepGraph, registryFileHashes, yankedButAllowedModules);
   }
 }

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/IndexRegistry.java

@@ -21,6 +21,7 @@
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.bazel.bzlmod.Version.ParseException;
 import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
 import com.google.devtools.build.lib.bazel.repository.downloader.Checksum.MissingChecksumException;
@@ -44,7 +45,6 @@
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Optional;
-import java.util.function.Predicate;
 
 /**
  * Represents a Bazel module registry that serves a list of module metadata from a static HTTP
@@ -69,6 +69,7 @@ public enum KnownFileHashesMode {
   private final Map<String, String> clientEnv;
   private final Gson gson;
   private final ImmutableMap<String, Optional<Checksum>> knownFileHashes;
+  private final ImmutableSet<ModuleKey> yankedButAllowedModules;
   private final KnownFileHashesMode knownFileHashesMode;
   private volatile Optional<BazelRegistryJson> bazelRegistryJson;
   private volatile StoredEventHandler bazelRegistryJsonEvents;
@@ -80,7 +81,8 @@ public IndexRegistry(
       DownloadManager downloadManager,
       Map<String, String> clientEnv,
       ImmutableMap<String, Optional<Checksum>> knownFileHashes,
-      KnownFileHashesMode knownFileHashesMode) {
+      KnownFileHashesMode knownFileHashesMode,
+      ImmutableSet<ModuleKey> yankedButAllowedModules) {
     this.uri = uri;
     this.downloadManager = downloadManager;
     this.clientEnv = clientEnv;
@@ -90,6 +92,7 @@ public IndexRegistry(
             .create();
     this.knownFileHashes = knownFileHashes;
     this.knownFileHashesMode = knownFileHashesMode;
+    this.yankedButAllowedModules = yankedButAllowedModules;
   }
 
   @Override
@@ -436,21 +439,16 @@ public Optional<ImmutableMap<Version, String>> getYankedVersions(
   }
 
   @Override
-  public boolean shouldFetchYankedVersions(
-      ModuleKey selectedModuleKey, Predicate<String> fileHashIsKnown) {
+  public boolean shouldFetchYankedVersions(ModuleKey selectedModuleKey) {
     // If the source.json hash is known, this module has been selected before when selection
-    // succeeded, which means that
-    // either:
+    // succeeded, which means that either:
     // * it wasn't yanked at that point in time and any successful selection since then has not seen
-    // a higher module
-    //   version, or
+    //   a higher module version, or
     // * it was yanked at that point in time, but explicitly allowed via
-    // BZLMOD_ALLOW_YANKED_VERSIONS.
-    // In both cases, we don't fetch yanked versions.
-    // TODO: Should we store whether we are in the second case and refetch yanked versions if the
-    // environment variable
-    //  changes?
-    return !fileHashIsKnown.test(getSourceJsonUrl(selectedModuleKey));
+    //   BZLMOD_ALLOW_YANKED_VERSIONS.
+    // In the first case, we don't fetch yanked versions.
+    return yankedButAllowedModules.contains(selectedModuleKey)
+        || !knownFileHashes.containsKey(getSourceJsonUrl(selectedModuleKey));
   }
 
   /** Represents fields available in {@code metadata.json} for each module. */

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/Registry.java

@@ -20,7 +20,6 @@
 import com.google.devtools.build.skyframe.SkyValue;
 import java.io.IOException;
 import java.util.Optional;
-import java.util.function.Predicate;
 
 /** A database where module metadata is stored. */
 public interface Registry extends SkyValue {
@@ -52,8 +51,7 @@ Optional<ImmutableMap<Version, String>> getYankedVersions(
 
   /**
    * Returns whether the (mutable) yanked versions should be fetched from the registry for the given
-   * module contained in the post-selection dependency graph, given the knowledge of certain
-   * registry file hashes contained in the lockfile.
+   * module contained in the post-selection dependency graph.
    */
-  boolean shouldFetchYankedVersions(ModuleKey selectedModuleKey, Predicate<String> fileHashIsKnown);
+  boolean shouldFetchYankedVersions(ModuleKey selectedModuleKey);
 }

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RegistryFactory.java

@@ -16,6 +16,7 @@
 package com.google.devtools.build.lib.bazel.bzlmod;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions;
 import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
 import java.net.URISyntaxException;
@@ -32,6 +33,7 @@ public interface RegistryFactory {
   Registry createRegistry(
       String url,
       ImmutableMap<String, Optional<Checksum>> fileHashes,
-      RepositoryOptions.LockfileMode lockfileMode)
+      RepositoryOptions.LockfileMode lockfileMode,
+      ImmutableSet<ModuleKey> yankedButAllowedModules)
       throws URISyntaxException;
 }

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RegistryFactoryImpl.java

@@ -16,6 +16,7 @@
 package com.google.devtools.build.lib.bazel.bzlmod;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.bazel.bzlmod.IndexRegistry.KnownFileHashesMode;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions.LockfileMode;
 import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
@@ -41,7 +42,8 @@ public RegistryFactoryImpl(
   public Registry createRegistry(
       String url,
       ImmutableMap<String, Optional<Checksum>> knownFileHashes,
-      LockfileMode lockfileMode)
+      LockfileMode lockfileMode,
+      ImmutableSet<ModuleKey> yankedButAllowedModules)
       throws URISyntaxException {
     URI uri = new URI(url);
     if (uri.getScheme() == null) {
@@ -71,6 +73,7 @@ public Registry createRegistry(
         downloadManager,
         clientEnvironmentSupplier.get(),
         knownFileHashes,
-        knownFileHashesMode);
+        knownFileHashesMode,
+        yankedButAllowedModules);
   }
 }

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RegistryFunction.java

@@ -51,7 +51,8 @@ public SkyValue compute(SkyKey skyKey, Environment env)
       return registryFactory.createRegistry(
           key.getUrl().replace("%workspace%", workspaceRoot.getPathString()),
           lockfile.getRegistryFileHashes(),
-          lockfileMode);
+          lockfileMode,
+          lockfile.getYankedButAllowedModules());
     } catch (URISyntaxException e) {
       throw new RegistryException(
           ExternalDepsException.withCauseAndMessage(

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/YankedVersionsFunction.java

@@ -46,8 +46,7 @@ public SkyValue compute(SkyKey skyKey, Environment env) throws InterruptedExcept
       return null;
     }
 
-    if (!registry.shouldFetchYankedVersions(
-        key.getModuleKey(), lockfile.getRegistryFileHashes()::containsKey)) {
+    if (!registry.shouldFetchYankedVersions(key.getModuleKey())) {
       return YankedVersionsValue.create(Optional.empty());
     }
 

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/YankedVersionsUtil.java

@@ -59,25 +59,6 @@ static Optional<ImmutableSet<ModuleKey>> parseAllowedYankedVersions(
     return Optional.of(allowedYankedVersionBuilder.build());
   }
 
-  static Optional<String> getYankedInfo(
-      ModuleKey key,
-      YankedVersionsValue yankedVersionsValue,
-      Optional<ImmutableSet<ModuleKey>> allowedYankedVersions) {
-    return yankedVersionsValue
-        .yankedVersions()
-        .flatMap(
-            yankedVersions -> {
-              String yankedInfo = yankedVersions.get(key.getVersion());
-              if (yankedInfo != null
-                  && allowedYankedVersions.isPresent()
-                  && !allowedYankedVersions.get().contains(key)) {
-                return Optional.of(yankedInfo);
-              } else {
-                return Optional.empty();
-              }
-            });
-  }
-
   /**
    * Parse of a comma-separated list of module version(s) of the form '<module name>@<version>' or
    * 'all' from the string. Returns true if 'all' is present, otherwise returns false.

src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BazelDepGraphFunctionTest.java

@@ -366,7 +366,11 @@ public void setDepGraph(ImmutableMap<ModuleKey, Module> depGraph) {
     @Override
     @Nullable
     public SkyValue compute(SkyKey skyKey, Environment env) {
-      return BazelModuleResolutionValue.create(depGraph, ImmutableMap.of(), ImmutableMap.of());
+      return BazelModuleResolutionValue.create(
+          depGraph,
+          ImmutableMap.of(),
+          ImmutableMap.of(),
+          state.discoverAndSelectResult.yankedButAllowedModules);
     }
   }
 }