diff --git a/src/model/ModelArtifactProvider.ts b/src/model/ModelArtifactProvider.ts index 49b0735c46..a696bab531 100644 --- a/src/model/ModelArtifactProvider.ts +++ b/src/model/ModelArtifactProvider.ts @@ -40,7 +40,6 @@ import { GgufHeader } from "./GgufHeader.js"; import { readSafetensorsHeader } from "./SafetensorsHeader.js"; import { ModelFormatFailure, - rejectPickle, MODEL_REPORT_BYTES, modelCapacityRemediation, modelCapacityReportedLimits, @@ -254,7 +253,6 @@ export class ModelArtifactProvider implements ModelArtifactPort { ); const reader = new ModelHeaderReader(file, Number(initial.size), signal); const first = await reader.take(4); - if (first.equals(Buffer.from([0x50, 0x4b, 3, 4]))) rejectPickle(); const format = first.equals(Buffer.from("GGUF")) ? "gguf" : "safetensors"; const partial: ModelArtifactPartial = { kind: "model-artifact-partial", diff --git a/src/model/SafetensorsHeader.ts b/src/model/SafetensorsHeader.ts index f6407dc8f9..5250c15479 100644 --- a/src/model/SafetensorsHeader.ts +++ b/src/model/SafetensorsHeader.ts @@ -215,14 +215,19 @@ export const readSafetensorsHeader = async ( // 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(); + // ZIP's four-byte signature is also a valid little-endian header length. + const zip = first.equals(Buffer.from([0x50, 0x4b, 3, 4])); + if ((pickle || zip) && reader.size < 9) rejectPickle(); const second = await reader.take(4); reader.excludeSafetensorsLengthPrefix(); const length = BigInt(first.readUInt32LE()) + (BigInt(second.readUInt32LE()) << 32n); 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))) + if ( + (pickle || zip) && + (brace !== 0x7b || length > BigInt(MODEL_HEADER_BYTES)) + ) rejectPickle(); if (brace !== 0x7b) rejectModel("File does not contain a GGUF or safetensors header"); diff --git a/tests/acceptance/modelSafetensorsZipCollision.test.ts b/tests/acceptance/modelSafetensorsZipCollision.test.ts new file mode 100644 index 0000000000..a3b35d5cca --- /dev/null +++ b/tests/acceptance/modelSafetensorsZipCollision.test.ts @@ -0,0 +1,60 @@ +import { open, writeFile } from "node:fs/promises"; +import { join } from "node:path"; +import { expect } from "vitest"; +import { createTestTempDirectory } from "../fixtures/temporaryDirectory.js"; +import { cliTest } from "../support/cli/cliFixture.js"; + +cliTest( + "classifies a ZIP-signature safetensors length before applying its representation limit", + async ({ cli }) => { + const root = await createTestTempDirectory("rea-model-zip-collision-"); + const length = 67_324_752; + const path = join(root, "model.safetensors"); + const prefix = Buffer.alloc(8); + prefix.writeBigUInt64LE(BigInt(length)); + const header = Buffer.from( + JSON.stringify({ + weight: { dtype: "U8", shape: [1], data_offsets: [0, 1] }, + }), + ); + const padding = Buffer.alloc(1024 * 1024, 0x20); + const file = await open(path, "wx"); + try { + await file.writeFile(prefix); + await file.writeFile(header); + for (let remaining = length - header.length; remaining > 0;) { + const count = Math.min(remaining, padding.length); + await file.writeFile(padding.subarray(0, count)); + remaining -= count; + } + await file.writeFile(Buffer.from([0x5a])); + } finally { + await file.close(); + } + const result = await cli.run({ + arguments: ["inspect-model-artifact", path, "--json"], + }); + expect(result.exitCode).toBe(1); + expect(result.json).toMatchObject({ + code: "resource_constraint", + category: "resource_constraint", + details: { + resource: "memory", + reported_limits: { maximum_representation_bytes: 33_554_432 }, + }, + }); + const zip = join(root, "archive.zip"); + await writeFile( + zip, + Buffer.from( + "504b0304140000000000000021008bc5848d05000000050000000a00000074656e736f722e747874696e657274504b01021403140000000000000021008bc5848d05000000050000000a000000000000000000000080010000000074656e736f722e747874504b05060000000001000100380000002d0000000000", + "hex", + ), + ); + const unsupported = await cli.run({ + arguments: ["inspect-model-artifact", zip, "--json"], + }); + expect(unsupported.exitCode).toBe(1); + expect(unsupported.json).toMatchObject({ code: "unsupported_target" }); + }, +); diff --git a/tests/boundary/model/modelArtifact.test.ts b/tests/boundary/model/modelArtifact.test.ts index 374f1c895c..be6aef9dfc 100644 --- a/tests/boundary/model/modelArtifact.test.ts +++ b/tests/boundary/model/modelArtifact.test.ts @@ -312,7 +312,7 @@ it("accepts standardized GGUF metadata keys containing numeric hierarchy compone }); }); -it.each([99_999_993, 100_000_000])( +it.each([67_324_752, 67_324_760, 99_999_993, 100_000_000])( "admits safetensors declared header length %s before its independent representation limit", async (length) => { const bytes = Buffer.alloc(9); @@ -1173,6 +1173,12 @@ it.each([640, 896, 1152, 1408, 66_176])( it.each([ // Protocol 2 pickle of {"a": 1}: PROTO 2, EMPTY_DICT, BINPUT 0, ... + ["a ZIP local header", Buffer.from("504b0304140000000000", "hex")], + ["a short ZIP signature", Buffer.from("504b0304", "hex")], + [ + "a ZIP signature followed by a brace and an inadmissible length", + Buffer.from("504b0304140000007b", "hex"), + ], ["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")],