1 files changed, 40 insertions, 3 deletions
diff --git a/src/main/java/com/google/devtools/build/lib/vfs/DigestHashFunction.java b/src/main/java/com/google/devtools/build/lib/vfs/DigestHashFunction.java
index 123e4b5a9c..e059ba832d 100644
--- a/src/main/java/com/google/devtools/build/lib/vfs/DigestHashFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/vfs/DigestHashFunction.java
@@ -18,6 +18,7 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.hash.HashFunction;
 import com.google.common.hash.Hashing;
+import com.google.devtools.build.lib.vfs.DigestHashFunction.DigestLength.DigestLengthImpl;
 import com.google.devtools.common.options.Converter;
 import com.google.devtools.common.options.OptionsParsingException;
 import java.security.MessageDigest;
@@ -39,6 +40,31 @@ public class DigestHashFunction {
   // This map must be declared first to make sure that calls to register() have it ready.
   private static final HashMap<String, DigestHashFunction> hashFunctionRegistry = new HashMap<>();
 
+  /** Describes the length of a digest. */
+  public interface DigestLength {
+    /** Returns the length of a digest by inspecting its bytes. Used for variable-length digests. */
+    default int getDigestLength(byte[] bytes, int offset) {
+      return getDigestMaximumLength();
+    }
+
+    /** Returns the maximum length a digest can turn into. */
+    int getDigestMaximumLength();
+
+    /** Default implementation that simply returns a fixed length. */
+    class DigestLengthImpl implements DigestLength {
+      private final int length;
+
+      DigestLengthImpl(HashFunction hashFunction) {
+        this.length = hashFunction.bits() / 8;
+      }
+
+      @Override
+      public int getDigestMaximumLength() {
+        return length;
+      }
+    }
+  }
+
   public static final DigestHashFunction MD5 = register(Hashing.md5(), "MD5");
   public static final DigestHashFunction SHA1 = register(Hashing.sha1(), "SHA-1", "SHA1");
   public static final DigestHashFunction SHA256 = register(Hashing.sha256(), "SHA-256", "SHA256");
@@ -47,17 +73,24 @@ public class DigestHashFunction {
   private static boolean defaultHasBeenSet = false;
 
   private final HashFunction hashFunction;
+  private final DigestLength digestLength;
   private final String name;
   private final MessageDigest messageDigestPrototype;
   private final boolean messageDigestPrototypeSupportsClone;
 
-  private DigestHashFunction(HashFunction hashFunction, String name) {
+  private DigestHashFunction(HashFunction hashFunction, DigestLength digestLength, String name) {
     this.hashFunction = hashFunction;
+    this.digestLength = digestLength;
     this.name = name;
     this.messageDigestPrototype = getMessageDigestInstance();
     this.messageDigestPrototypeSupportsClone = supportsClone(messageDigestPrototype);
   }
 
+  public static DigestHashFunction register(
+      HashFunction hash, String hashName, String... altNames) {
+    return register(hash, new DigestLengthImpl(hash), hashName, altNames);
+  }
+
   /**
    * Creates a new DigestHashFunction that is registered to be recognized by its name in {@link
    * DigestFunctionConverter}.
@@ -70,7 +103,7 @@ public class DigestHashFunction {
    * @throws IllegalArgumentException if the name is already registered.
    */
   public static DigestHashFunction register(
-      HashFunction hash, String hashName, String... altNames) {
+      HashFunction hash, DigestLength digestLength, String hashName, String... altNames) {
     try {
       MessageDigest.getInstance(hashName);
     } catch (NoSuchAlgorithmException e) {
@@ -80,7 +113,7 @@ public class DigestHashFunction {
           e);
     }
 
-    DigestHashFunction hashFunction = new DigestHashFunction(hash, hashName);
+    DigestHashFunction hashFunction = new DigestHashFunction(hash, digestLength, hashName);
     List<String> names = ImmutableList.<String>builder().add(hashName).add(altNames).build();
     synchronized (hashFunctionRegistry) {
       for (String name : names) {
@@ -177,6 +210,10 @@ public class DigestHashFunction {
     }
   }
 
+  public DigestLength getDigestLength() {
+    return digestLength;
+  }
+
   public boolean isValidDigest(byte[] digest) {
     // TODO(b/109764197): Remove this check to accept variable-length hashes.
     return digest != null && digest.length * 8 == hashFunction.bits();