[api] Uses sha-256 to avoid security warning (#2495)

deepjavalibrary · Mar 31, 2023 · 6aedd50 · 6aedd50
1 parent 212bc14
commit 6aedd50
Show file tree

Hide file tree

Showing 8 changed files with 43 additions and 16 deletions.
diff --git a/api/src/main/java/ai/djl/repository/AbstractRepository.java b/api/src/main/java/ai/djl/repository/AbstractRepository.java
@@ -30,7 +30,6 @@
 import java.io.UnsupportedEncodingException;
 import java.net.URI;
 import java.net.URLDecoder;
-import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -264,15 +263,6 @@ private void untar(InputStream is, Path dir, boolean gzip) throws IOException {
         }
     }
 
-    protected static String md5hash(String input) {
-        try {
-            MessageDigest md = MessageDigest.getInstance("MD5");
-            return Hex.toHexString(md.digest(input.getBytes(StandardCharsets.UTF_8)));
-        } catch (NoSuchAlgorithmException e) {
-            throw new AssertionError("MD5 algorithm not found.", e);
-        }
-    }
-
     private static Map<String, String> parseQueryString(URI uri) {
         try {
             Map<String, String> map = new ConcurrentHashMap<>();

diff --git a/api/src/main/java/ai/djl/repository/JarRepository.java b/api/src/main/java/ai/djl/repository/JarRepository.java
@@ -15,6 +15,7 @@
 import ai.djl.Application;
 import ai.djl.repository.zoo.DefaultModelZoo;
 import ai.djl.util.Progress;
+import ai.djl.util.Utils;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -122,7 +123,8 @@ private synchronized Metadata getMetadata() {
         metadata = new Metadata.MatchAllMetadata();
         metadata.setArtifactId(artifactId);
         metadata.setArtifacts(Collections.singletonList(artifact));
-        String hash = md5hash(queryString == null ? uri.toString() : uri.toString() + queryString);
+        String hash =
+                Utils.hash(queryString == null ? uri.toString() : uri.toString() + queryString);
         MRL mrl = model(Application.UNDEFINED, DefaultModelZoo.GROUP_ID, hash);
         metadata.setRepositoryUri(mrl.toURI());
 

diff --git a/api/src/main/java/ai/djl/repository/SimpleRepository.java b/api/src/main/java/ai/djl/repository/SimpleRepository.java
@@ -16,6 +16,7 @@
 import ai.djl.repository.Artifact.Item;
 import ai.djl.repository.zoo.DefaultModelZoo;
 import ai.djl.util.Progress;
+import ai.djl.util.Utils;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -192,7 +193,7 @@ private synchronized Metadata getMetadata() throws IOException {
             files.put(artifactId, item);
             artifact.setFiles(files);
 
-            String hash = md5hash(uri);
+            String hash = Utils.hash(uri);
             MRL mrl = model(Application.UNDEFINED, DefaultModelZoo.GROUP_ID, hash);
             metadata.setRepositoryUri(mrl.toURI());
         } else {

diff --git a/api/src/main/java/ai/djl/repository/SimpleUrlRepository.java b/api/src/main/java/ai/djl/repository/SimpleUrlRepository.java
@@ -15,6 +15,7 @@
 import ai.djl.Application;
 import ai.djl.repository.zoo.DefaultModelZoo;
 import ai.djl.util.Progress;
+import ai.djl.util.Utils;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -124,7 +125,7 @@ private synchronized Metadata getMetadata() throws IOException {
         metadata = new Metadata.MatchAllMetadata();
         metadata.setArtifactId(artifactId);
         metadata.setArtifacts(Collections.singletonList(artifact));
-        String hash = md5hash(uri.toString());
+        String hash = Utils.hash(uri.toString());
         MRL mrl = model(Application.UNDEFINED, DefaultModelZoo.GROUP_ID, hash);
         metadata.setRepositoryUri(mrl.toURI());
         return metadata;

diff --git a/api/src/main/java/ai/djl/util/Hex.java b/api/src/main/java/ai/djl/util/Hex.java
@@ -28,12 +28,25 @@ private Hex() {}
      * @return the converted hex String
      */
     public static String toHexString(byte[] block) {
+        return toHexString(block, 0, block.length);
+    }
+
+    /**
+     * Converts a byte array to a hex string.
+     *
+     * @param block the bytes to convert
+     * @param start the start position (inclusive) of the array
+     * @param end the end position (exclusive) of the array
+     * @return the converted hex String
+     */
+    public static String toHexString(byte[] block, int start, int end) {
         if (block == null) {
             return null;
         }
 
         StringBuilder buf = new StringBuilder();
-        for (byte aBlock : block) {
+        for (int i = start; i < end; ++i) {
+            byte aBlock = block[i];
             int high = ((aBlock & 0xf0) >> 4);
             int low = (aBlock & 0x0f);
 

diff --git a/api/src/main/java/ai/djl/util/Utils.java b/api/src/main/java/ai/djl/util/Utils.java
@@ -29,6 +29,8 @@
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.nio.file.StandardCopyOption;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
@@ -460,4 +462,20 @@ public static InputStream openUrl(URL url) throws IOException {
         }
         return new BufferedInputStream(url.openStream());
     }
+
+    /**
+     * Returns a hash of a string.
+     *
+     * @param input the input string
+     * @return a 20 bytes hash of the input stream in hex format
+     */
+    public static String hash(String input) {
+        try {
+            MessageDigest md = MessageDigest.getInstance("SHA-256");
+            byte[] buf = md.digest(input.getBytes(StandardCharsets.UTF_8));
+            return Hex.toHexString(buf, 0, 20);
+        } catch (NoSuchAlgorithmException e) {
+            throw new AssertionError("SHA256 algorithm not found.", e);
+        }
+    }
 }
diff --git a/extensions/aws-ai/src/main/java/ai/djl/aws/s3/S3Repository.java b/extensions/aws-ai/src/main/java/ai/djl/aws/s3/S3Repository.java
@@ -21,6 +21,7 @@
 import ai.djl.repository.Repository;
 import ai.djl.repository.zoo.DefaultModelZoo;
 import ai.djl.util.Progress;
+import ai.djl.util.Utils;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -165,7 +166,7 @@ private synchronized Metadata getMetadata() throws IOException {
             }
 
             metadata = new Metadata.MatchAllMetadata();
-            String hash = md5hash("s3://" + bucket + '/' + prefix);
+            String hash = Utils.hash("s3://" + bucket + '/' + prefix);
             MRL mrl = model(Application.UNDEFINED, DefaultModelZoo.GROUP_ID, hash);
             metadata.setRepositoryUri(mrl.toURI());
             metadata.setArtifactId(artifactId);

diff --git a/extensions/hadoop/src/main/java/ai/djl/hadoop/hdfs/HdfsRepository.java b/extensions/hadoop/src/main/java/ai/djl/hadoop/hdfs/HdfsRepository.java
@@ -21,6 +21,7 @@
 import ai.djl.repository.Repository;
 import ai.djl.repository.zoo.DefaultModelZoo;
 import ai.djl.util.Progress;
+import ai.djl.util.Utils;
 
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FileStatus;
@@ -161,7 +162,7 @@ private synchronized Metadata getMetadata() throws IOException {
         }
 
         metadata = new Metadata.MatchAllMetadata();
-        String hash = md5hash(uri.resolve(prefix).toString());
+        String hash = Utils.hash(uri.resolve(prefix).toString());
         MRL mrl = model(Application.UNDEFINED, DefaultModelZoo.GROUP_ID, hash);
         metadata.setRepositoryUri(mrl.toURI());
         metadata.setArtifactId(artifactId);