diff --git a/h3_safetensors.c b/h3_safetensors.c index 60ab177a..dc81ea98 100644 --- a/h3_safetensors.c +++ b/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; diff --git a/tests/test_h3.c b/tests/test_h3.c index 3f66fa32..34a71118 100644 --- a/tests/test_h3.c +++ b/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) {