diff --git a/packages/cli/src/utils/download.test.ts b/packages/cli/src/utils/download.test.ts index 84f0f9305d..3e8b362ce5 100644 --- a/packages/cli/src/utils/download.test.ts +++ b/packages/cli/src/utils/download.test.ts @@ -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 = {}, + body?: string | ((url: string) => string), +): typeof httpsGet { + return ((_url: string, callback: (response: IncomingMessage) => void) => { + const response = new PassThrough() as PassThrough & { + statusCode: number; + headers: Record; + }; + 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("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", + ); + }); +}); diff --git a/packages/cli/src/utils/download.ts b/packages/cli/src/utils/download.ts index 6e5a06b9ed..ae6d34af78 100644 --- a/packages/cli/src/utils/download.ts +++ b/packages/cli/src/utils/download.ts @@ -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 { - 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`)); });