fix(cli): follow every redirect a host may answer with (#3148)

* fix(cli): follow every redirect a host may answer with

downloadFile handled 301 and 302 and passed the Location header straight back as a request target. Hosts answer with relative locations far more often than that assumed, and 303, 307 and 308 are all reachable, so a CDN handoff failed on a URL that was never a URL.

Locations now resolve against the URL that sent them, the code set covers all five, and a hop cap ends a redirect loop rather than recursing.

Keeps mains idle-response test, which asserts the request timeout fires and clears the partial file. An earlier version of this branch replaced the file wholesale and lost it, leaving downloadFile with no test that calls it at all.

* fix(cli): bound and isolate model downloads
This commit is contained in:
Miguel Ángel
2026-08-09 21:49:43 -07:00
committed by GitHub
parent c96b30c717
commit db3de4c1bf
2 changed files with 203 additions and 30 deletions
+107 -2
View File
@@ -1,12 +1,12 @@
import { EventEmitter } from "node:events";
import { existsSync, mkdtempSync, rmSync } from "node:fs";
import { existsSync, mkdtempSync, readFileSync, readdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { PassThrough } from "node:stream";
import { get as httpsGet } from "node:https";
import type { ClientRequest, IncomingMessage } from "node:http";
import { afterEach, describe, expect, it, vi } from "vitest";
import { downloadFile } from "./download.js";
import { downloadFile, REDIRECT_CODES, redirectTarget } from "./download.js";
const { unlinkSyncMock } = vi.hoisted(() => ({
unlinkSyncMock: vi.fn(),
@@ -22,12 +22,61 @@ vi.mock("node:fs", async (importOriginal) => {
const mockGet = vi.mocked(httpsGet);
const tempDirs: string[] = [];
function httpsResponse(
statusCode: number,
headers: Record<string, string> = {},
body?: string | ((url: string) => string),
): typeof httpsGet {
return ((_url: string, callback: (response: IncomingMessage) => void) => {
const response = new PassThrough() as PassThrough & {
statusCode: number;
headers: Record<string, string>;
};
response.statusCode = statusCode;
response.headers = headers;
const request = new EventEmitter() as ClientRequest;
request.setTimeout = vi.fn();
callback(response as unknown as IncomingMessage);
if (body !== undefined) {
queueMicrotask(() => response.end(typeof body === "function" ? body(_url) : body));
}
return request;
}) as typeof httpsGet;
}
afterEach(() => {
vi.clearAllMocks();
for (const dir of tempDirs.splice(0)) rmSync(dir, { recursive: true, force: true });
});
describe("downloadFile", () => {
it("keeps concurrent partial downloads separate", async () => {
mockGet.mockImplementation(httpsResponse(200, {}, (url) => url));
const dir = mkdtempSync(join(tmpdir(), "hyperframes-download-"));
tempDirs.push(dir);
const dest = join(dir, "model.onnx");
const first = "https://example.test/first";
const second = "https://example.test/second";
await Promise.all([downloadFile(first, dest), downloadFile(second, dest)]);
expect([first, second]).toContain(readFileSync(dest, "utf-8"));
});
it("rejects a response that exceeds its byte limit", async () => {
mockGet.mockImplementation(httpsResponse(200, {}, "too large"));
const dir = mkdtempSync(join(tmpdir(), "hyperframes-download-"));
tempDirs.push(dir);
const dest = join(dir, "model.onnx");
await expect(
downloadFile("https://example.test/model.onnx", dest, { maxBytes: 3 }),
).rejects.toThrow("Download exceeded 3 bytes");
expect(existsSync(dest)).toBe(false);
expect(readdirSync(dir)).toEqual([]);
});
it("rejects an idle response and removes the partial file", async () => {
const actualFs = await vi.importActual<typeof import("node:fs")>("node:fs");
let responseClosed = false;
@@ -87,3 +136,59 @@ describe("downloadFile", () => {
expect(existsSync(dest)).toBe(false);
});
});
describe("redirect handling", () => {
it("rejects a redirect that the HTTPS client cannot follow", async () => {
mockGet
.mockImplementationOnce(
httpsResponse(307, { location: "http://cdn.example.test/model.onnx" }),
)
.mockImplementationOnce(() => {
throw new TypeError('Protocol "http:" not supported');
});
const dir = mkdtempSync(join(tmpdir(), "hyperframes-download-"));
tempDirs.push(dir);
await expect(
downloadFile("https://example.test/model.onnx", join(dir, "model.onnx")),
).rejects.toThrow('Protocol "http:" not supported');
});
it("follows every redirect a host may answer with", () => {
// 307 is the one that broke: HuggingFace answers the tokenizer with it,
// and the original set stopped at 302, so the download fell through to the
// not-200 branch and hung.
for (const code of [301, 302, 303, 307, 308]) {
expect(REDIRECT_CODES.has(code)).toBe(true);
}
});
it("does not treat a success or an error as a redirect", () => {
for (const code of [200, 204, 404, 500]) {
expect(REDIRECT_CODES.has(code)).toBe(false);
}
});
it("resolves a relative location against the url that sent it", () => {
// The actual failure: handing this path back as a request target is not a
// valid URL, so nothing was ever fetched.
expect(
redirectTarget(
"/api/resolve-cache/models/x/tokenizer.json",
"https://huggingface.co/x/resolve/main/tokenizer.json",
),
).toBe("https://huggingface.co/api/resolve-cache/models/x/tokenizer.json");
});
it("keeps an absolute location as it is", () => {
expect(redirectTarget("https://cdn.example.com/blob", "https://huggingface.co/a/b")).toBe(
"https://cdn.example.com/blob",
);
});
it("resolves a sibling-relative location", () => {
expect(redirectTarget("other.json", "https://example.com/a/b/file.json")).toBe(
"https://example.com/a/b/other.json",
);
});
});
+96 -28
View File
@@ -1,6 +1,8 @@
import { randomUUID } from "node:crypto";
import { createWriteStream, renameSync, unlinkSync } from "node:fs";
import { get as httpsGet } from "node:https";
import type { IncomingMessage } from "node:http";
import type { ClientRequest, IncomingMessage } from "node:http";
import { Transform } from "node:stream";
import { pipeline } from "node:stream/promises";
const DEFAULT_DOWNLOAD_TIMEOUT_MS = 30_000;
@@ -8,8 +10,28 @@ const DEFAULT_DOWNLOAD_TIMEOUT_MS = 30_000;
export interface DownloadOptions {
/** Abort after this many milliseconds without network activity. */
timeoutMs?: number;
/** Reject before writing more than this many response bytes. */
maxBytes?: number;
}
/** Every redirect a host may reasonably answer with, not just the two we saw first. */
export const REDIRECT_CODES = new Set([301, 302, 303, 307, 308]);
/**
* Where a redirect actually points.
*
* Split out because this is the part that was wrong, and it is the part that
* can be checked without a socket: hosts answer with relative locations
* (`/api/resolve-cache/...`) far more often than the original code assumed, and
* handing that string back as a request target fails.
*/
export function redirectTarget(location: string, from: string): string {
return new URL(location, from).toString();
}
/** Enough hops for a CDN handoff, few enough that a redirect loop still ends. */
const MAX_REDIRECTS = 10;
function removePartialFile(path: string): void {
try {
unlinkSync(path);
@@ -18,51 +40,97 @@ function removePartialFile(path: string): void {
}
}
function enforceByteLimit(maxBytes: number): Transform {
let received = 0;
return new Transform({
transform(chunk: Buffer, _encoding, callback) {
received += chunk.byteLength;
callback(
received > maxBytes ? new Error(`Download exceeded ${maxBytes} bytes`) : null,
chunk,
);
},
});
}
/**
* Download a file from a URL, following redirects.
* Uses atomic write (download to .tmp, rename on success) to prevent
* corrupt partial files from persisting in the cache on interruption.
*
* Location headers are resolved against the URL that sent them, because they
* are frequently relative: a host answering 307 with `/api/resolve-cache/...`
* is normal, and passing that string back as a request target is not.
*
* The timeout is per request, so it re-arms on each hop of a redirect chain
* rather than budgeting the whole chain. A stalled socket is what it exists to
* catch, and without it the CLI waits forever.
*/
export function downloadFile(
url: string,
dest: string,
options: DownloadOptions = {},
): Promise<void> {
const tmp = `${dest}.tmp`;
const tmp = `${dest}.${process.pid}.${randomUUID()}.tmp`;
const timeoutMs = options.timeoutMs ?? DEFAULT_DOWNLOAD_TIMEOUT_MS;
const maxBytes = options.maxBytes;
return new Promise((resolve, reject) => {
const follow = (u: string) => {
const follow = (u: string, hops = 0) => {
let activeResponse: IncomingMessage | undefined;
let responsePipelineStarted = false;
let requestError: Error | undefined;
const request = httpsGet(u, (res) => {
activeResponse = res;
if (res.statusCode === 301 || res.statusCode === 302) {
const location = res.headers.location;
if (location) {
let request: ClientRequest;
try {
request = httpsGet(u, (res) => {
activeResponse = res;
if (res.statusCode && REDIRECT_CODES.has(res.statusCode)) {
const location = res.headers.location;
if (location) {
if (hops >= MAX_REDIRECTS) {
res.resume();
removePartialFile(tmp);
reject(
new Error(`Download failed: more than ${MAX_REDIRECTS} redirects from ${url}`),
);
return;
}
res.resume();
try {
follow(redirectTarget(location, u), hops + 1);
} catch (error) {
removePartialFile(tmp);
reject(error);
}
return;
}
}
if (res.statusCode !== 200) {
res.resume();
follow(location);
removePartialFile(tmp);
reject(new Error(`Download failed: HTTP ${res.statusCode}`));
return;
}
}
if (res.statusCode !== 200) {
res.resume();
removePartialFile(tmp);
reject(new Error(`Download failed: HTTP ${res.statusCode}`));
return;
}
const file = createWriteStream(tmp);
responsePipelineStarted = true;
pipeline(res, file)
.then(() => {
renameSync(tmp, dest);
resolve();
})
.catch((err) => {
removePartialFile(tmp);
reject(requestError ?? err);
});
});
const file = createWriteStream(tmp);
responsePipelineStarted = true;
const transfer =
maxBytes === undefined
? pipeline(res, file)
: pipeline(res, enforceByteLimit(maxBytes), file);
transfer
.then(() => {
renameSync(tmp, dest);
resolve();
})
.catch((err) => {
removePartialFile(tmp);
reject(requestError ?? err);
});
});
} catch (error) {
removePartialFile(tmp);
reject(error);
return;
}
request.setTimeout(timeoutMs, () => {
request.destroy(new Error(`Download timed out after ${timeoutMs}ms`));
});