Reject overlapping safetensors data offsets

h3_safetensors.c

@@ -356,6 +356,34 @@ static int h3_append_tensor(h3_st_header *header, h3_st_tensor tensor,
     return 1;
 }
 
+static int h3_tensor_offset_compare(const void *left, const void *right) {
+    const h3_st_tensor *a = left;
+    const h3_st_tensor *b = right;
+    if (a->data_begin != b->data_begin)
+        return a->data_begin < b->data_begin ? -1 : 1;
+    if (a->data_end != b->data_end)
+        return a->data_end < b->data_end ? -1 : 1;
+    return 0;
+}
+
+static int h3_validate_offsets(h3_json_cursor *cursor, h3_st_header *header,
+                               uint64_t data_size) {
+    if (header->tensor_count > 1) {
+        qsort(header->tensors, header->tensor_count, sizeof(*header->tensors),
+              h3_tensor_offset_compare);
+    }
+    uint64_t expected = 0;
+    for (size_t index = 0; index < header->tensor_count; index++) {
+        if (header->tensors[index].data_begin != expected) {
+            return h3_json_fail(cursor, "tensor data offsets are not contiguous");
+        }
+        expected = header->tensors[index].data_end;
+    }
+    if (expected != data_size)
+        return h3_json_fail(cursor, "tensor data offsets do not cover file");
+    return 1;
+}
+
 static uint64_t h3_u64_le(const unsigned char bytes[8]) {
     uint64_t value = 0;
     for (unsigned index = 0; index < 8; index++) {
@@ -450,6 +478,8 @@ int h3_st_read_header(const char *path, h3_st_header *header,
         h3_json_fail(&cursor, "trailing data in safetensors header");
         goto fail;
     }
+    if (!h3_validate_offsets(&cursor, header, file_size - 8 - header_size))
+        goto fail;
     free(json);
     return 1;
 

tests/test_h3.c

@@ -258,8 +258,8 @@ static void test_safetensors(void) {
     int descriptor = mkstemp(path);
     CHECK(descriptor >= 0);
     const char header_json[] =
-        "{\"x\":{\"dtype\":\"F32\",\"shape\":[2,3],\"data_offsets\":[0,24]},"
-        "\"scalar\":{\"dtype\":\"BF16\",\"shape\":[],\"data_offsets\":[24,26]}}";
+        "{\"scalar\":{\"dtype\":\"BF16\",\"shape\":[],\"data_offsets\":[24,26]},"
+        "\"x\":{\"dtype\":\"F32\",\"shape\":[2,3],\"data_offsets\":[0,24]}}";
     uint64_t length = sizeof(header_json) - 1;
     unsigned char prefix[8];
     for (unsigned index = 0; index < 8; index++) prefix[index] = (unsigned char)(length >> (8 * index));
@@ -292,6 +292,23 @@ static void test_safetensors(void) {
     CHECK(h3_st_tensor_elements(scalar) == 1);
     h3_st_free_header(&header);
     CHECK(unlink(path) == 0);
+
+    char overlap_path[] = "/tmp/h3_safetensors_overlap_XXXXXX";
+    descriptor = mkstemp(overlap_path);
+    CHECK(descriptor >= 0);
+    const char overlap_json[] =
+        "{\"x\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[0,4]},"
+        "\"y\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[2,6]}}";
+    length = sizeof(overlap_json) - 1;
+    for (unsigned index = 0; index < 8; index++)
+        prefix[index] = (unsigned char)(length >> (8 * index));
+    unsigned char overlap_payload[6] = {0};
+    write_all(descriptor, prefix, sizeof(prefix));
+    write_all(descriptor, overlap_json, (size_t)length);
+    write_all(descriptor, overlap_payload, sizeof(overlap_payload));
+    CHECK(close(descriptor) == 0);
+    CHECK(!h3_st_read_header(overlap_path, &header, error, sizeof(error)));
+    CHECK(unlink(overlap_path) == 0);
 }
 
 static void test_rng_and_solver(void) {