Merge pull request #135 from c2an1/fix-issue-#134

Fix issue #134
tjake · Dec 27, 2024 · a496949 · a496949
2 parents 6da3b7d + 6c0329e
commit a496949
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/SafeTensorSupport.java b/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/SafeTensorSupport.java
@@ -47,8 +47,19 @@ public class SafeTensorSupport {
     private static final MapType metadataTypeReference = om.getTypeFactory().constructMapType(Map.class, String.class, String.class);
 
     public static Map<String, TensorInfo> readTensorInfoMap(ByteBuffer buf, Optional<Map<String, String>> saveMetadata) {
+        final long MAX_HEADER_LENGTH = 1024 * 1024 * 1024; // 1 GB
         buf = buf.order(ByteOrder.LITTLE_ENDIAN);
         long headerLength = buf.getLong();
+
+        // headerLength is negative
+        if (headerLength < 0) {
+            throw new IllegalArgumentException("Header length cannot be negative: " + headerLength);
+        }
+        // headerLength exceeds the maximum allowed length MAX_HEADER_LENGTH
+        if (headerLength > MAX_HEADER_LENGTH) {
+            throw new IllegalArgumentException(String.format("Header length %d exceeds the maximum allowed length %d.", headerLength, MAX_HEADER_LENGTH));
+        }
+
         byte[] header = new byte[Ints.checkedCast(headerLength)];
         buf.get(header);