diff --git a/src/model/ModelArtifactProvider.ts b/src/model/ModelArtifactProvider.ts index 8e5efa2541..cc0d0a4936 100644 --- a/src/model/ModelArtifactProvider.ts +++ b/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") { diff --git a/src/model/ModelFormat.ts b/src/model/ModelFormat.ts index 2e8d36956a..560415fb18 100644 --- a/src/model/ModelFormat.ts +++ b/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. diff --git a/src/model/SafetensorsHeader.ts b/src/model/SafetensorsHeader.ts index 601bdd1b8d..9fc24b2113 100644 --- a/src/model/SafetensorsHeader.ts +++ b/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( diff --git a/tests/boundary/model/modelArtifact.test.ts b/tests/boundary/model/modelArtifact.test.ts index 72896a7996..9fc4ae1574 100644 --- a/tests/boundary/model/modelArtifact.test.ts +++ b/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);