metrale_core/
safetensors.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The one place that turns a safetensors `data_offsets` pair into a byte span.
4//!
5//! Owner: metrale-core.
6//! Invariants:
7//! - A returned `TensorSpan` satisfies `abs_offset + len <= file_len`, computed
8//!   without overflow; every other input is an error.
9//!
10//! A checkpoint is third-party data, so every number in its header is
11//! untrusted until checked. The header declares each tensor as
12//! `"data_offsets": [start, end]` relative to the data section, and the naive
13//! `end - start` is a `u64` subtraction that wraps on a crafted or truncated
14//! file: a reversed pair yields a length near `u64::MAX`.
15//!
16//! Both header parsers call this: `metrale_model_weights::fast_weights`
17//! (`fast_weights/header.rs`, whose spans become `pread` windows) and
18//! `metrale_storage::weight_peer` (`weight_peer/shard.rs`, which builds the
19//! RDMA manifest).
20
21use anyhow::{Context, Result, bail};
22use serde_json::Value;
23
24/// 2026-09-25: A validated tensor byte span: an absolute file offset and a
25/// length that is known to fit inside the file it came from.
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub struct TensorSpan {
28    /// 2026-09-25: Absolute byte offset in the shard file where the tensor starts.
29    pub abs_offset: u64,
30    pub len: u64,
31}
32
33/// 2026-09-25: Validate one tensor's `data_offsets` against the file that declared it.
34///
35/// `offsets` is the raw `data_offsets` JSON value, `data_start` is
36/// `8 + header_size` (where the data section begins), and `file_len` is the
37/// shard's size on disk.
38///
39/// Rejects, in this order: a value that is not an array of two `u64`s, a
40/// reversed pair (`end < start`, the underflow), an absolute start or end
41/// that overflows `u64`, and a span that runs past end-of-file. Callers get a
42/// named error instead of a wrapped length.
43pub fn tensor_span(
44    tensor: &str,
45    offsets: &Value,
46    data_start: u64,
47    file_len: u64,
48) -> Result<TensorSpan> {
49    let arr = offsets
50        .as_array()
51        .with_context(|| format!("tensor {tensor}: data_offsets is not an array"))?;
52    if arr.len() != 2 {
53        bail!(
54            "tensor {tensor}: data_offsets has {} entries, expected 2",
55            arr.len()
56        );
57    }
58    let rel_start = arr[0]
59        .as_u64()
60        .with_context(|| format!("tensor {tensor}: bad data_offsets[0]"))?;
61    let rel_end = arr[1]
62        .as_u64()
63        .with_context(|| format!("tensor {tensor}: bad data_offsets[1]"))?;
64    // 2026-09-25: The underflow. `rel_end - rel_start` wraps to ~u64::MAX on a reversed
65    // pair; in release builds that is silent.
66    if rel_end < rel_start {
67        bail!("tensor {tensor}: data_offsets [{rel_start}, {rel_end}] end precedes start");
68    }
69    let len = rel_end - rel_start;
70    // 2026-09-25: `data_start` is bounded by the 64 MiB header cap both callers
71    // enforce, but `rel_start` is not, so the sum still needs checking.
72    let abs_offset = data_start
73        .checked_add(rel_start)
74        .with_context(|| format!("tensor {tensor}: data_offsets[0] {rel_start} overflows"))?;
75    let abs_end = abs_offset
76        .checked_add(len)
77        .with_context(|| format!("tensor {tensor}: span {abs_offset}+{len} overflows"))?;
78    if abs_end > file_len {
79        bail!(
80            "tensor {tensor}: spans bytes {abs_offset}..{abs_end} but the shard is only \
81             {file_len} bytes (truncated or corrupt checkpoint)"
82        );
83    }
84    Ok(TensorSpan { abs_offset, len })
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90    use serde_json::json;
91
92    /// 2026-09-25: A well-formed pair keeps its exact span, converted to absolute.
93    #[test]
94    fn accepts_well_formed_span() {
95        let s = tensor_span("w", &json!([0, 128]), 72, 4096).unwrap();
96        assert_eq!(
97            s,
98            TensorSpan {
99                abs_offset: 72,
100                len: 128
101            }
102        );
103        let s = tensor_span("w", &json!([128, 256]), 72, 4096).unwrap();
104        assert_eq!(
105            s,
106            TensorSpan {
107                abs_offset: 200,
108                len: 128
109            }
110        );
111    }
112
113    /// 2026-09-25: A zero-length tensor is accepted.
114    #[test]
115    fn accepts_empty_tensor() {
116        let s = tensor_span("w", &json!([64, 64]), 8, 4096).unwrap();
117        assert_eq!(s.len, 0);
118    }
119
120    /// 2026-09-25: A tensor may end exactly at EOF.
121    #[test]
122    fn accepts_span_ending_exactly_at_eof() {
123        let s = tensor_span("w", &json!([0, 4024]), 72, 4096).unwrap();
124        assert_eq!(s.len, 4024);
125    }
126
127    /// 2026-09-25: `rel_end - rel_start` on a reversed pair would wrap to
128    /// ~u64::MAX in release builds; it is rejected instead.
129    #[test]
130    fn rejects_reversed_offsets_instead_of_underflowing() {
131        let err = tensor_span("evil", &json!([4096, 0]), 8, 65536).unwrap_err();
132        let msg = err.to_string();
133        assert!(msg.contains("evil"), "{msg}");
134        assert!(msg.contains("end precedes start"), "{msg}");
135        // 2026-09-25: Sanity-check that the unchecked form really does wrap, so this test
136        // is guarding a live hazard and not a hypothetical one.
137        assert_eq!(0u64.wrapping_sub(4096), u64::MAX - 4095);
138    }
139
140    /// 2026-09-25: A truncated shard: the header still advertises the full tensor.
141    #[test]
142    fn rejects_span_past_end_of_file() {
143        let err = tensor_span("w", &json!([0, 1_000_000]), 72, 4096).unwrap_err();
144        assert!(err.to_string().contains("truncated or corrupt"), "{err}");
145    }
146
147    /// 2026-09-25: A huge `rel_start` must not wrap when added to `data_start`.
148    #[test]
149    fn rejects_offset_that_overflows_u64() {
150        let err = tensor_span("w", &json!([u64::MAX, u64::MAX]), 72, 4096).unwrap_err();
151        assert!(
152            err.to_string().contains("data_offsets[0]") && err.to_string().contains("overflows"),
153            "{err}"
154        );
155    }
156
157    /// 2026-09-25: A valid absolute start can still overflow when its nonzero length is added.
158    #[test]
159    fn rejects_span_end_that_overflows_u64() {
160        let err = tensor_span("w", &json!([u64::MAX - 1, u64::MAX]), 1, u64::MAX).unwrap_err();
161        assert!(
162            err.to_string().contains("span") && err.to_string().contains("overflows"),
163            "{err}"
164        );
165    }
166
167    /// 2026-09-25: Malformed `data_offsets` shapes are named, not indexed into blindly.
168    #[test]
169    fn rejects_malformed_offsets() {
170        let cases = [
171            (json!("nope"), "not an array"),
172            (json!([0]), "1 entries"),
173            (json!([0, 1, 2]), "3 entries"),
174            (json!([-1, 16]), "data_offsets[0]"),
175            (json!([0.5, 16]), "data_offsets[0]"),
176            (json!([0, -1]), "data_offsets[1]"),
177            (json!([0, 16.5]), "data_offsets[1]"),
178            (json!([0, "16"]), "data_offsets[1]"),
179        ];
180        for (offsets, cause) in cases {
181            let err = tensor_span("w", &offsets, 8, 4096).unwrap_err();
182            let msg = err.to_string();
183            assert!(msg.contains("tensor w"), "{msg}");
184            assert!(msg.contains(cause), "expected {cause} in: {msg}");
185        }
186    }
187}
metrale_core/safetensors.rs

metrale_core/
safetensors.rs
