From bbf6a7fbbab412b551ef72bc4ba281ac4670aa5b Mon Sep 17 00:00:00 2001 From: Draco Glasser Date: Sat, 22 Aug 2026 17:44:07 +0800 Subject: [PATCH] Reject ambiguous safetensors payload layouts The safetensors format relies on tensor data offsets fully describing the byte buffer: offsets are relative to the payload, must be ordered without gaps or overlaps after sorting by offset, and the final end must match the payload size. h3 already checked each tensor shape against its own byte range, but it accepted layouts where two tensor ranges overlapped or left bytes unindexed. This validates the global offset layout once the header has been parsed and adds a regression with an out-of-order valid header plus overlapping tensor ranges. Constraint: Safetensors README requires the byte buffer to be entirely indexed and hole-free to prevent polyglot payloads. Rejected: Pairwise overlap-only scan | would still accept trailing unindexed bytes. Confidence: high Scope-risk: narrow Directive: Keep this check at header parse time so all safetensors consumers share the same trust-boundary validation. Tested: make test Tested: ASan/UBSan h3_tests Not-tested: Full model-weight parity tests requiring released fixtures not present locally --- h3_safetensors.c | 30 ++++++++++++++++++++++++++++++ tests/test_h3.c | 21 +++++++++++++++++++-- 2 files changed, 49 insertions(+), 2 deletions(-) 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) {