From 22638b9b7883b58ec2ce98f675bb9dc9b41826e1 Mon Sep 17 00:00:00 2001 From: Michael D'Angelo Date: Fri, 10 Apr 2026 11:35:11 -0700 Subject: [PATCH 1/2] fix: mark malformed safetensors inconclusive --- CHANGELOG.md | 1 + modelaudit/scanners/safetensors_scanner.py | 177 +++++++++++++++++---- tests/scanners/test_safetensors_scanner.py | 79 ++++++++- 3 files changed, 224 insertions(+), 33 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fc1c12a9c..0a1f0934d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Bug Fixes +- mark malformed SafeTensors framing inconclusive instead of allowing clean exits - avoid CoreML nested parse failures on bounded-read truncation - flag TensorFlow `LoadLibrary` and `LoadLibraryV2` graph ops as dangerous native-library loading - detect split CNTK native-user-function and native-library references diff --git a/modelaudit/scanners/safetensors_scanner.py b/modelaudit/scanners/safetensors_scanner.py index ccbbdac5f..78ba7ebc6 100644 --- a/modelaudit/scanners/safetensors_scanner.py +++ b/modelaudit/scanners/safetensors_scanner.py @@ -10,7 +10,7 @@ from modelaudit.detectors.suspicious_symbols import SUSPICIOUS_METADATA_PATTERNS -from .base import BaseScanner, IssueSeverity, ScanResult +from .base import INCONCLUSIVE_SCAN_OUTCOME, BaseScanner, IssueSeverity, ScanResult # Map SafeTensors dtypes to byte sizes for integrity checking _DTYPE_SIZES = { @@ -31,6 +31,9 @@ "U64": 8, } MAX_HEADER_BYTES = 16 * 1024 * 1024 +SAFETENSORS_HEADER_INCONCLUSIVE_REASON = "safetensors_header_validation_failed" +SAFETENSORS_STRUCTURE_INCONCLUSIVE_REASON = "safetensors_structure_validation_failed" +SAFETENSORS_HEADER_LIMIT_INCONCLUSIVE_REASON = "safetensors_header_size_limit_exceeded" class SafeTensorsScanner(BaseScanner): @@ -40,6 +43,26 @@ class SafeTensorsScanner(BaseScanner): description = "Scans SafeTensors model files for integrity issues" supported_extensions: ClassVar[list[str]] = [".safetensors"] + @staticmethod + def _mark_inconclusive(result: ScanResult, reason: str) -> None: + """Mark malformed safetensors framing as an explicit inconclusive scan.""" + result.metadata["analysis_incomplete"] = True + result.metadata["scan_outcome"] = INCONCLUSIVE_SCAN_OUTCOME + + reasons = result.metadata.get("scan_outcome_reasons") + if not isinstance(reasons, list): + reasons = [] + result.metadata["scan_outcome_reasons"] = reasons + if reason not in reasons: + reasons.append(reason) + + @staticmethod + def _is_valid_shape(shape: Any) -> bool: + """Return True when shape is a safetensors-compatible list of non-negative ints.""" + return isinstance(shape, list) and all( + isinstance(dim, int) and not isinstance(dim, bool) and dim >= 0 for dim in shape + ) + @classmethod def can_handle(cls, path: str) -> bool: """Check if this scanner can handle the given path.""" @@ -70,6 +93,7 @@ def scan(self, path: str) -> ScanResult: result = self._create_result() file_size = self.get_file_size(path) result.metadata["file_size"] = file_size + structural_validation_failed = False # Add file integrity check for compliance self.add_file_integrity_check(path, result) @@ -87,6 +111,7 @@ def scan(self, path: str) -> ScanResult: location=path, details={"bytes_read": len(header_len_bytes), "required": 8}, ) + self._mark_inconclusive(result, SAFETENSORS_HEADER_INCONCLUSIVE_REASON) result.finish(success=False) return result @@ -101,6 +126,7 @@ def scan(self, path: str) -> ScanResult: location=path, details={"header_len": header_len, "max_allowed": file_size - 8}, ) + self._mark_inconclusive(result, SAFETENSORS_HEADER_INCONCLUSIVE_REASON) result.finish(success=False) return result else: @@ -128,6 +154,7 @@ def scan(self, path: str) -> ScanResult: ), ) result.metadata["analysis_incomplete"] = True + self._mark_inconclusive(result, SAFETENSORS_HEADER_LIMIT_INCONCLUSIVE_REASON) result.bytes_scanned = file_size result.finish(success=True) return result @@ -150,6 +177,7 @@ def scan(self, path: str) -> ScanResult: location=path, details={"bytes_read": len(header_bytes), "expected": header_len}, ) + self._mark_inconclusive(result, SAFETENSORS_HEADER_INCONCLUSIVE_REASON) result.finish(success=False) return result @@ -161,6 +189,7 @@ def scan(self, path: str) -> ScanResult: severity=IssueSeverity.INFO, location=path, ) + self._mark_inconclusive(result, SAFETENSORS_HEADER_INCONCLUSIVE_REASON) result.finish(success=False) return result else: @@ -183,6 +212,7 @@ def scan(self, path: str) -> ScanResult: details={"exception": str(e), "exception_type": type(e).__name__}, why="SafeTensors header contained invalid JSON.", ) + self._mark_inconclusive(result, SAFETENSORS_HEADER_INCONCLUSIVE_REASON) result.finish(success=False) return result @@ -213,13 +243,40 @@ def scan(self, path: str) -> ScanResult: location=path, details={"tensor": name, "actual_type": type(info).__name__, "expected_type": "dict"}, ) + self._mark_inconclusive(result, SAFETENSORS_STRUCTURE_INCONCLUSIVE_REASON) + structural_validation_failed = True continue - begin, end = info.get("data_offsets", [0, 0]) + raw_offsets = info.get("data_offsets") dtype = info.get("dtype") shape = info.get("shape", []) - if not isinstance(begin, int) or not isinstance(end, int): + if not isinstance(raw_offsets, list) or len(raw_offsets) != 2: + result.add_check( + name="Tensor Offset Structure Validation", + passed=False, + message=f"Invalid data_offsets structure for {name}", + severity=IssueSeverity.INFO, + location=path, + details={ + "tensor": name, + "actual_type": type(raw_offsets).__name__, + "expected_type": "list", + "expected_length": 2, + }, + ) + self._mark_inconclusive(result, SAFETENSORS_STRUCTURE_INCONCLUSIVE_REASON) + structural_validation_failed = True + continue + + begin, end = raw_offsets + + if ( + not isinstance(begin, int) + or isinstance(begin, bool) + or not isinstance(end, int) + or isinstance(end, bool) + ): result.add_check( name="Tensor Offset Type Validation", passed=False, @@ -232,6 +289,8 @@ def scan(self, path: str) -> ScanResult: "end_type": type(end).__name__, }, ) + self._mark_inconclusive(result, SAFETENSORS_STRUCTURE_INCONCLUSIVE_REASON) + structural_validation_failed = True continue if begin < 0 or end <= begin or end > data_size: @@ -256,32 +315,82 @@ def scan(self, path: str) -> ScanResult: offsets.append((begin, end)) # Validate dtype/shape size + if not isinstance(dtype, str) or dtype not in _DTYPE_SIZES: + result.add_check( + name="Tensor Dtype Validation", + passed=False, + message=f"Invalid dtype for tensor {name}", + severity=IssueSeverity.INFO, + location=path, + details={ + "tensor": name, + "dtype": dtype, + "actual_type": type(dtype).__name__, + }, + ) + self._mark_inconclusive(result, SAFETENSORS_STRUCTURE_INCONCLUSIVE_REASON) + structural_validation_failed = True + continue + + if not self._is_valid_shape(shape): + result.add_check( + name="Tensor Shape Validation", + passed=False, + message=f"Invalid shape for tensor {name}", + severity=IssueSeverity.INFO, + location=path, + details={ + "tensor": name, + "shape": shape, + "actual_type": type(shape).__name__, + }, + ) + self._mark_inconclusive(result, SAFETENSORS_STRUCTURE_INCONCLUSIVE_REASON) + structural_validation_failed = True + continue + expected_size = self._expected_size(dtype, shape) - if expected_size is not None: - if expected_size != end - begin: - result.add_check( - name="Tensor Size Consistency Check", - passed=False, - message=f"Size mismatch for tensor {name}", - severity=IssueSeverity.CRITICAL, - location=path, - details={ - "tensor": name, - "expected_size": expected_size, - "actual_size": end - begin, - }, - ) - else: - result.add_check( - name="Tensor Size Consistency Check", - passed=True, - message=f"Tensor {name} size matches dtype/shape", - location=path, - details={ - "tensor": name, - "size": expected_size, - }, - ) + if expected_size is None: + result.add_check( + name="Tensor Size Computation Check", + passed=False, + message=f"Unable to compute expected size for tensor {name}", + severity=IssueSeverity.INFO, + location=path, + details={ + "tensor": name, + "dtype": dtype, + "shape": shape, + }, + ) + self._mark_inconclusive(result, SAFETENSORS_STRUCTURE_INCONCLUSIVE_REASON) + structural_validation_failed = True + continue + + if expected_size != end - begin: + result.add_check( + name="Tensor Size Consistency Check", + passed=False, + message=f"Size mismatch for tensor {name}", + severity=IssueSeverity.CRITICAL, + location=path, + details={ + "tensor": name, + "expected_size": expected_size, + "actual_size": end - begin, + }, + ) + else: + result.add_check( + name="Tensor Size Consistency Check", + passed=True, + message=f"Tensor {name} size matches dtype/shape", + location=path, + details={ + "tensor": name, + "size": expected_size, + }, + ) # Check offset continuity offsets.sort(key=lambda x: x[0]) @@ -298,6 +407,8 @@ def scan(self, path: str) -> ScanResult: location=path, details={"gap_at": begin, "expected": last_end}, ) + self._mark_inconclusive(result, SAFETENSORS_STRUCTURE_INCONCLUSIVE_REASON) + structural_validation_failed = True break last_end = end @@ -320,6 +431,8 @@ def scan(self, path: str) -> ScanResult: location=path, details={"last_offset": last_end, "data_size": data_size}, ) + self._mark_inconclusive(result, SAFETENSORS_STRUCTURE_INCONCLUSIVE_REASON) + structural_validation_failed = True # Check metadata metadata = header.get("__metadata__", {}) @@ -388,18 +501,20 @@ def scan(self, path: str) -> ScanResult: result.finish(success=False) return result - result.finish(success=not result.has_errors) + result.finish(success=not result.has_errors and not structural_validation_failed) return result @staticmethod - def _expected_size(dtype: str | None, shape: list[int]) -> int | None: + def _expected_size(dtype: str | None, shape: Any) -> int | None: """Return expected tensor byte size from dtype and shape.""" if dtype not in _DTYPE_SIZES: return None + if not isinstance(shape, list): + return None size = _DTYPE_SIZES[dtype] total = 1 for dim in shape: - if not isinstance(dim, int) or dim < 0: + if not isinstance(dim, int) or isinstance(dim, bool) or dim < 0: return None total *= dim return total * size diff --git a/tests/scanners/test_safetensors_scanner.py b/tests/scanners/test_safetensors_scanner.py index 3e0314046..87c51d063 100644 --- a/tests/scanners/test_safetensors_scanner.py +++ b/tests/scanners/test_safetensors_scanner.py @@ -14,7 +14,8 @@ from safetensors.numpy import save_file -from modelaudit.scanners.base import CheckStatus +from modelaudit.core import determine_exit_code, scan_model_directory_or_file +from modelaudit.scanners.base import INCONCLUSIVE_SCAN_OUTCOME, CheckStatus, IssueSeverity from modelaudit.scanners.safetensors_scanner import SafeTensorsScanner @@ -43,6 +44,11 @@ def create_safetensors_with_dtype_size_mismatch(path: Path, dtype: str) -> None: path.write_bytes(struct.pack(" None: + header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8") + path.write_bytes(struct.pack(" None: file_path = tmp_path / "model.safetensors" create_safetensors_file(file_path) @@ -94,6 +100,8 @@ def test_oversized_header_triggers_limit_check(tmp_path: Path) -> None: assert "exceeds maximum allowed size" in header_limit_check.message assert result.success is True assert result.metadata["analysis_incomplete"] is True + assert result.metadata["scan_outcome"] == INCONCLUSIVE_SCAN_OUTCOME + assert "safetensors_header_size_limit_exceeded" in result.metadata["scan_outcome_reasons"] assert result.bytes_scanned == file_path.stat().st_size @@ -117,6 +125,7 @@ def track_analyze(metadata: dict[str, object], result: object, path: str) -> Non assert header_limit_check is not None assert header_limit_check.status.value == "failed" assert result.metadata["analysis_incomplete"] is True + assert result.metadata["scan_outcome"] == INCONCLUSIVE_SCAN_OUTCOME assert result.success is True @@ -229,6 +238,71 @@ def test_bad_offsets(tmp_path: Path) -> None: assert any("offset" in issue.message.lower() for issue in result.issues) +def test_unclaimed_safetensors_data_is_inconclusive_not_clean(tmp_path: Path) -> None: + file_path = tmp_path / "trailing.safetensors" + header = {"t": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]}} + write_raw_safetensors(file_path, header, b"\x00" * 8) + + scanner = SafeTensorsScanner() + result = scanner.scan(str(file_path)) + + assert result.success is False + assert result.has_errors is False + assert result.metadata["scan_outcome"] == INCONCLUSIVE_SCAN_OUTCOME + assert "safetensors_structure_validation_failed" in result.metadata["scan_outcome_reasons"] + assert not any(issue.severity == IssueSeverity.CRITICAL for issue in result.issues) + assert any( + check.name == "Tensor Data Coverage Check" and check.status == CheckStatus.FAILED for check in result.checks + ) + + +def test_unclaimed_safetensors_data_returns_exit2(tmp_path: Path) -> None: + file_path = tmp_path / "trailing.safetensors" + header = {"t": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]}} + write_raw_safetensors(file_path, header, b"\x00" * 8) + + result = scan_model_directory_or_file(str(file_path)) + + assert determine_exit_code(result) == 2 + assert not any(issue.severity == IssueSeverity.CRITICAL for issue in result.issues) + + +def test_malformed_data_offsets_are_inconclusive_not_scanner_crash(tmp_path: Path) -> None: + file_path = tmp_path / "bad_offsets_shape.safetensors" + header = {"t": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4, 8]}} + write_raw_safetensors(file_path, header, b"\x00" * 8) + + scanner = SafeTensorsScanner() + result = scanner.scan(str(file_path)) + + assert result.success is False + assert result.has_errors is False + assert result.metadata["scan_outcome"] == INCONCLUSIVE_SCAN_OUTCOME + assert "safetensors_structure_validation_failed" in result.metadata["scan_outcome_reasons"] + assert any( + check.name == "Tensor Offset Structure Validation" and check.status == CheckStatus.FAILED + for check in result.checks + ) + assert not any("Error scanning SafeTensors file" in issue.message for issue in result.issues) + + +def test_safetensors_security_finding_takes_precedence_over_inconclusive_structure(tmp_path: Path) -> None: + file_path = tmp_path / "malicious_metadata_trailing.safetensors" + header = { + "__metadata__": {"description": ""}, + "t": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]}, + } + write_raw_safetensors(file_path, header, b"\x00" * 8) + + direct = SafeTensorsScanner().scan(str(file_path)) + aggregate = scan_model_directory_or_file(str(file_path)) + + assert direct.has_errors is True + assert direct.metadata["scan_outcome"] == INCONCLUSIVE_SCAN_OUTCOME + assert any(issue.severity == IssueSeverity.CRITICAL for issue in direct.issues) + assert determine_exit_code(aggregate) == 1 + + @pytest.mark.parametrize( ("dtype", "expected_size"), [("BOOL", 4), ("BF16", 8), ("F8_E4M3", 4), ("F8_E5M2", 4), ("F16", 8), ("F32", 16), ("F64", 32)], @@ -275,7 +349,7 @@ def test_deeply_nested_header(tmp_path: Path) -> None: scanner = SafeTensorsScanner() result = scanner.scan(str(file_path)) - assert result.has_errors + assert result.has_errors or result.metadata.get("scan_outcome") == INCONCLUSIVE_SCAN_OUTCOME # Check that either RecursionError was caught OR the header was marked as invalid/deeply nested # Also check for generic JSON error since deeply nested JSON might fail differently # Include tensor validation errors as acceptable since deeply nested but valid JSON @@ -286,6 +360,7 @@ def test_deeply_nested_header(tmp_path: Path) -> None: or "recursion" in check.message.lower() or "invalid json" in check.message.lower() or "offsets out of bounds" in check.message.lower() # Acceptable for this test + or "invalid data_offsets structure" in check.message.lower() for check in result.checks ) From 017cc9fc903865a5825034ccd562f48911bc7a19 Mon Sep 17 00:00:00 2001 From: Michael D'Angelo Date: Fri, 10 Apr 2026 16:05:06 -0700 Subject: [PATCH 2/2] fix: handle safetensors unicode header errors --- modelaudit/scanners/safetensors_scanner.py | 3 +- tests/scanners/test_safetensors_scanner.py | 34 ++++++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/modelaudit/scanners/safetensors_scanner.py b/modelaudit/scanners/safetensors_scanner.py index 78ba7ebc6..e827fc5ba 100644 --- a/modelaudit/scanners/safetensors_scanner.py +++ b/modelaudit/scanners/safetensors_scanner.py @@ -202,7 +202,7 @@ def scan(self, path: str) -> ScanResult: try: header = json.loads(header_bytes.decode("utf-8")) - except json.JSONDecodeError as e: + except (UnicodeDecodeError, json.JSONDecodeError) as e: result.add_check( name="SafeTensors JSON Parse", passed=False, @@ -215,7 +215,6 @@ def scan(self, path: str) -> ScanResult: self._mark_inconclusive(result, SAFETENSORS_HEADER_INCONCLUSIVE_REASON) result.finish(success=False) return result - tensor_names = [k for k in header if k != "__metadata__"] result.metadata["tensor_count"] = len(tensor_names) result.metadata["tensors"] = tensor_names diff --git a/tests/scanners/test_safetensors_scanner.py b/tests/scanners/test_safetensors_scanner.py index 87c51d063..30336e01d 100644 --- a/tests/scanners/test_safetensors_scanner.py +++ b/tests/scanners/test_safetensors_scanner.py @@ -49,6 +49,10 @@ def write_raw_safetensors(path: Path, header: dict[str, Any], data: bytes) -> No path.write_bytes(struct.pack(" None: + path.write_bytes(struct.pack(" None: file_path = tmp_path / "model.safetensors" create_safetensors_file(file_path) @@ -214,6 +218,36 @@ def test_corrupted_header(tmp_path: Path) -> None: assert any("json" in msg or "header" in msg or "invalid" in msg or "corrupt" in msg for msg in all_messages) +def test_non_object_header_is_inconclusive_not_clean(tmp_path: Path) -> None: + file_path = tmp_path / "array_header.safetensors" + write_raw_safetensors_header(file_path, b"[]") + + direct = SafeTensorsScanner().scan(str(file_path)) + + assert direct.success is False + assert direct.has_errors is False + assert direct.metadata["scan_outcome"] == INCONCLUSIVE_SCAN_OUTCOME + assert "safetensors_header_validation_failed" in direct.metadata["scan_outcome_reasons"] + assert any( + check.name == "Header Format Validation" and check.status == CheckStatus.FAILED for check in direct.checks + ) + assert not any(issue.severity == IssueSeverity.CRITICAL for issue in direct.issues) + + +def test_invalid_utf8_header_is_inconclusive_not_scanner_crash(tmp_path: Path) -> None: + file_path = tmp_path / "invalid_utf8_header.safetensors" + write_raw_safetensors_header(file_path, b"{\xff}") + + direct = SafeTensorsScanner().scan(str(file_path)) + + assert direct.success is False + assert direct.has_errors is False + assert direct.metadata["scan_outcome"] == INCONCLUSIVE_SCAN_OUTCOME + assert "safetensors_header_validation_failed" in direct.metadata["scan_outcome_reasons"] + assert any(check.name == "SafeTensors JSON Parse" and check.status == CheckStatus.FAILED for check in direct.checks) + assert not any(issue.severity == IssueSeverity.CRITICAL for issue in direct.issues) + + def test_bad_offsets(tmp_path: Path) -> None: file_path = tmp_path / "model.safetensors" create_safetensors_file(file_path)