fix(model): recognize safetensors ZIP signature collisions (#7816)

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",

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");

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" });
+  },
+);

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")],