in safetensors/src/tensor.rs [1307:1347]
fn test_json_attack() {
let mut tensors = HashMap::new();
let dtype = Dtype::F32;
let shape = vec![2, 2];
let data_offsets = (0, 16);
for i in 0..10 {
tensors.insert(
format!("weight_{i}"),
TensorInfo {
dtype,
shape: shape.clone(),
data_offsets,
},
);
}
let metadata = HashMetadata {
metadata: None,
tensors,
};
let serialized = serde_json::to_string(&metadata).unwrap();
let serialized = serialized.as_bytes();
let n = serialized.len();
let filename = "out.safetensors";
let mut f = std::io::BufWriter::new(std::fs::File::create(filename).unwrap());
f.write_all(n.to_le_bytes().as_ref()).unwrap();
f.write_all(serialized).unwrap();
f.write_all(b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0").unwrap();
f.flush().unwrap();
let reloaded = std::fs::read(filename).unwrap();
match SafeTensors::deserialize(&reloaded) {
Err(SafeTensorError::InvalidOffset(_)) => {
// Yes we have the correct error, name of the tensor is random though
}
Err(err) => panic!("Unexpected error {err:?}"),
Ok(_) => panic!("This should not be able to be deserialized"),
}
}