fix(model): admit safetensors header lengths that resemble pickle markers (#5918)

A safetensors length prefix such as 640 encodes as 80 02 00 ..., which the
provider read as a pickle PROTO marker and refused before the safetensors
parser ran. The parser that owns the length prefix now makes the call: a
pickle-like prefix is safetensors only when byte 8 is '{' and the declared
length is within the safetensors header ceiling. Real pickles stay refused
and are never deserialized.

Fixes #5895

src/model/ModelArtifactProvider.ts

@@ -42,6 +42,7 @@ import { readSafetensorsHeader } from "./SafetensorsHeader.js";
 import {
   ModelFormatFailure,
   rejectModel,
+  rejectPickle,
   MODEL_REPORT_BYTES,
   modelCapacityRemediation,
   modelCapacityReportedLimits,
@@ -274,14 +275,7 @@ export class ModelArtifactProvider implements ModelArtifactPort {
       );
     const reader = new ModelHeaderReader(file, Number(initial.size), signal);
     const first = await reader.take(4);
-    if (
-      (first[0] === 0x80 && (first[1] ?? 0) >= 2 && (first[1] ?? 0) <= 5) ||
-      first.equals(Buffer.from([0x50, 0x4b, 3, 4]))
-    )
-      throw new ModelFormatFailure(
-        "unsupported",
-        "PyTorch pickle/ZIP model deserialization is unsupported; supply GGUF or safetensors instead",
-      );
+    if (first.equals(Buffer.from([0x50, 0x4b, 3, 4]))) rejectPickle();
     const format = first.equals(Buffer.from("GGUF")) ? "gguf" : "safetensors";
     let partial: ModelArtifactPartial | undefined;
     if (format === "gguf") {

src/model/ModelFormat.ts

@@ -35,6 +35,13 @@ export function rejectModel(message: string): never {
   throw new ModelFormatFailure("format", message);
 }
 
+export function rejectPickle(): never {
+  throw new ModelFormatFailure(
+    "unsupported",
+    "PyTorch pickle/ZIP model deserialization is unsupported; supply GGUF or safetensors instead",
+  );
+}
+
 // Shared header ceiling for GGUF and Safetensors. Safetensors counts declared
 // JSON bytes after the eight-byte length prefix; GGUF counts the complete
 // header. Retained rows and the serialized report have separate allowances.

src/model/SafetensorsHeader.ts

@@ -7,9 +7,11 @@ import type {
 } from "../domain/model/modelArtifact.js";
 import { ModelHeaderReader } from "./ModelHeaderReader.js";
 import {
+  MODEL_HEADER_BYTES,
   ModelFormatFailure,
   ModelReportBudget,
   rejectModel,
+  rejectPickle,
 } from "./ModelFormat.js";
 import { safetensorsBits } from "./ModelTensorTypes.js";
 
@@ -154,11 +156,19 @@ export const readSafetensorsHeader = async (
   tensors: ModelTensor[];
 }> => {
   const budget = new ModelReportBudget();
+  // Lengths such as 640 (80 02 ...) also begin like a pickle PROTO opcode.
+  const protocol = first[1] ?? 0;
+  const pickle = first[0] === 0x80 && protocol >= 2 && protocol <= 5;
+  if (pickle && reader.size < 9) rejectPickle();
   const second = await reader.take(4);
   reader.excludeSafetensorsLengthPrefix();
   const length =
     BigInt(first.readUInt32LE()) + (BigInt(second.readUInt32LE()) << 32n);
-  if (length === 0n || (await reader.take(1))[0] !== 0x7b)
+  const brace = length === 0n ? undefined : (await reader.take(1))[0];
+  // A brace after an admissible length is safetensors, even if truncated.
+  if (pickle && (brace !== 0x7b || length > BigInt(MODEL_HEADER_BYTES)))
+    rejectPickle();
+  if (brace !== 0x7b)
     rejectModel("File does not contain a GGUF or safetensors header");
   if (length > BigInt(Number.MAX_SAFE_INTEGER))
     throw new ModelFormatFailure(

tests/boundary/model/modelArtifact.test.ts

@@ -738,6 +738,80 @@ it.each(malformed)(
   },
 );
 
+it.each([640, 896, 1152, 1408, 66_176])(
+  "admits safetensors header length %s whose prefix bytes resemble a pickle protocol marker",
+  async (length) => {
+    const header = JSON.stringify({
+      weight: { dtype: "F32", shape: [1], data_offsets: [0, 4] },
+    });
+    const fixture = safetensorsFixture(
+      header.padEnd(length, " "),
+      Buffer.alloc(4),
+    );
+    const { provider, path } = await providerAt(fixture.bytes);
+    const result = await provider.inspect({ path, hash: "header" });
+    if (!result.ok) throw result.error;
+    expect(result.value).toMatchObject({ format: "safetensors" });
+    expect(result.value.tensors.map(({ name }) => name)).toEqual(["weight"]);
+  },
+);
+
+it.each([
+  // Protocol 2 pickle of {"a": 1}: PROTO 2, EMPTY_DICT, BINPUT 0, ...
+  ["a real pickle", Buffer.from("80027d710058010000006171014b01732e", "hex")],
+  // Protocol 2 pickle of {}: six bytes, shorter than a safetensors prefix.
+  ["an empty dict", Buffer.from("80027d71002e", "hex")],
+  // Protocol 4 FRAME pickle of {"a": 1}.
+  [
+    "a protocol 4 frame",
+    Buffer.from("8004950a000000000000007d948c0161944b01732e", "hex"),
+  ],
+  // The magic number torch.save writes first in its legacy format.
+  [
+    "the legacy torch.save magic",
+    Buffer.from("80028a0a6cfc9c46f9206aa850192e", "hex"),
+  ],
+  // Protocol 2 pickle of "x{": its brace lands at offset 8.
+  [
+    "a brace after an inadmissible length",
+    Buffer.from("80025802000000787b71002e", "hex"),
+  ],
+  [
+    "an in-file length without a JSON brace",
+    Buffer.concat([
+      Buffer.from("8002000000000000", "hex"),
+      Buffer.alloc(644, 0x2e),
+    ]),
+  ],
+] as const)(
+  "refuses pickle protocol bytes followed by %s",
+  async (_case, bytes) => {
+    const { provider, path } = await providerAt(bytes);
+    const result = await provider.inspect({ path, hash: "header" });
+    if (result.ok) throw new Error("Pickle admitted");
+    expect(result.error).toMatchObject({
+      _tag: "AnalysisUnsupportedTargetError",
+      reason: expect.stringContaining("PyTorch pickle"),
+    });
+  },
+);
+
+it("reports a truncated safetensors header whose length resembles a pickle marker as malformed", async () => {
+  const fixture = safetensorsFixture(
+    JSON.stringify({
+      weight: { dtype: "F32", shape: [1], data_offsets: [0, 4] },
+    }).padEnd(896, " "),
+    Buffer.alloc(4),
+  );
+  const { provider, path } = await providerAt(fixture.bytes.subarray(0, 500));
+  const result = await provider.inspect({ path, hash: "header" });
+  if (result.ok) throw new Error("Truncated model was accepted");
+  expect(result.error).toMatchObject({
+    _tag: "BinaryTargetError",
+    constraint: "file_format_rejected",
+  });
+});
+
 it("bounds expansion before allocation and distinguishes unsupported pickle, version and paths", async () => {
   const fixture = ggufFixture();
   fixture.bytes.writeBigUInt64LE(0xffffffffffffffffn, 8);