From 624053ee2f42e0e646c2452414e403abafc52781 Mon Sep 17 00:00:00 2001
From: chelsealong <chelsealong@126.com>
Date: Sun, 9 Aug 2026 01:46:06 +0000
Subject: [PATCH] Validate full safetensors data region is contiguous in
 no-mmap loader

Per-tensor ranges were checked individually but gaps, overlaps, a
non-zero first offset, and trailing bytes were still accepted. Sort
ranges by start and require the data region to be fully and
contiguously covered, matching safetensors' own validation on the
mmap path. Adds corruption tests for a gap, an overlap, and trailing
bytes.
---
 comfy/utils.py                                | 21 ++++++++-
 tests-unit/comfy_test/load_torch_file_test.py | 46 +++++++++++++++++++
 2 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/comfy/utils.py b/comfy/utils.py
index 175838547..0814f7ccd 100644
--- a/comfy/utils.py
+++ b/comfy/utils.py
@@ -135,6 +135,9 @@ def load_safetensors_no_mmap(ckpt, device):
             raise ValueError("Invalid safetensors header: header size exceeds the maximum allowed size")
         header = json.loads(f.read(header_size).decode("utf-8"))
         data_start = 8 + header_size
+        data_size = os.fstat(f.fileno()).st_size - data_start
+
+        tensors = []
         for name, info in header.items():
             if name == "__metadata__":
                 continue
@@ -142,9 +145,23 @@ def load_safetensors_no_mmap(ckpt, device):
             dtype = _TYPES[info["dtype"]]
             shape = info["shape"]
             expected_size = math.prod(shape) * dtype.itemsize
-            if end < start or end - start != expected_size:
+            if start < 0 or end < start or end - start != expected_size:
                 raise ValueError("Invalid safetensors header: tensor '{}' data range does not match its declared shape/dtype".format(name))
-            if expected_size == 0:
+            tensors.append((start, end, name, dtype, shape))
+
+        # The data region must be fully and contiguously indexed by the header,
+        # with no gaps, overlaps, or trailing bytes, matching the validation
+        # the safetensors library itself performs on the mmap path.
+        next_start = 0
+        for start, end, name, _dtype, _shape in sorted(tensors, key=lambda t: t[0]):
+            if start != next_start:
+                raise ValueError("Invalid safetensors header: tensor data ranges are not contiguous")
+            next_start = end
+        if next_start != data_size:
+            raise ValueError("Invalid safetensors header: tensor data does not cover the full file")
+
+        for start, end, name, dtype, shape in tensors:
+            if start == end:
                 sd[name] = torch.empty(shape, dtype=dtype, device=device)
                 continue
             f.seek(data_start + start)
diff --git a/tests-unit/comfy_test/load_torch_file_test.py b/tests-unit/comfy_test/load_torch_file_test.py
index eddd12a87..f64355bd7 100644
--- a/tests-unit/comfy_test/load_torch_file_test.py
+++ b/tests-unit/comfy_test/load_torch_file_test.py
@@ -48,3 +48,49 @@ def test_load_safetensors_no_mmap_rejects_corrupt_data_offsets(tmp_path):
 
     with pytest.raises(ValueError):
         comfy.utils.load_safetensors_no_mmap(str(path), torch.device("cpu"))
+
+
+def _write_safetensors_with_raw_offsets(path, header, data):
+    header_bytes = json.dumps(header).encode("utf-8")
+    with open(path, "wb") as f:
+        f.write(struct.pack("<Q", len(header_bytes)))
+        f.write(header_bytes)
+        f.write(data)
+
+
+def test_load_safetensors_no_mmap_rejects_gap_between_tensors(tmp_path):
+    # weight2 starts one byte after weight ends, leaving an unindexed gap.
+    header = {
+        "weight": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]},
+        "weight2": {"dtype": "F32", "shape": [1], "data_offsets": [5, 9]},
+    }
+    path = tmp_path / "gap.safetensors"
+    _write_safetensors_with_raw_offsets(path, header, b"\x00" * 9)
+
+    with pytest.raises(ValueError):
+        comfy.utils.load_safetensors_no_mmap(str(path), torch.device("cpu"))
+
+
+def test_load_safetensors_no_mmap_rejects_overlapping_tensors(tmp_path):
+    # weight2 starts before weight ends, so the ranges overlap.
+    header = {
+        "weight": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]},
+        "weight2": {"dtype": "F32", "shape": [1], "data_offsets": [2, 6]},
+    }
+    path = tmp_path / "overlap.safetensors"
+    _write_safetensors_with_raw_offsets(path, header, b"\x00" * 6)
+
+    with pytest.raises(ValueError):
+        comfy.utils.load_safetensors_no_mmap(str(path), torch.device("cpu"))
+
+
+def test_load_safetensors_no_mmap_rejects_trailing_bytes(tmp_path):
+    # The declared data region ends before the end of the file.
+    header = {
+        "weight": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]},
+    }
+    path = tmp_path / "trailing.safetensors"
+    _write_safetensors_with_raw_offsets(path, header, b"\x00" * 8)
+
+    with pytest.raises(ValueError):
+        comfy.utils.load_safetensors_no_mmap(str(path), torch.device("cpu"))