From 6c0329eeecd32a4a9c2c5c66209cd355952b923e Mon Sep 17 00:00:00 2001 From: c2an1 <12221068@zju.edu.cn> Date: Wed, 25 Dec 2024 12:08:59 +0800 Subject: [PATCH] Fix issue #134 --- .../tjake/jlama/safetensors/SafeTensorSupport.java | 11 +++++++++++ 1 file changed, 11 insertions(+) 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 index a9ad3c1..0af3ca2 100644 --- 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 readTensorInfoMap(ByteBuffer buf, Optional> 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);