Skip to content

Commit db3de4c

Browse files
miguel-heygenmiguel-heygen
authored
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
1 parent c96b30c commit db3de4c

2 files changed

Lines changed: 203 additions & 30 deletions

File tree

‎packages/cli/src/utils/download.test.ts‎

Lines changed: 107 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
import { EventEmitter } from "node:events";
2-
import { existsSync, mkdtempSync, rmSync } from "node:fs";
2+
import { existsSync, mkdtempSync, readFileSync, readdirSync, rmSync } from "node:fs";
33
import { tmpdir } from "node:os";
44
import { join } from "node:path";
55
import { PassThrough } from "node:stream";
66
import { get as httpsGet } from "node:https";
77
import type { ClientRequest, IncomingMessage } from "node:http";
88
import { afterEach, describe, expect, it, vi } from "vitest";
9-
import { downloadFile } from "./download.js";
9+
import { downloadFile, REDIRECT_CODES, redirectTarget } from "./download.js";
1010

1111
const { unlinkSyncMock } = vi.hoisted(() => ({
1212
unlinkSyncMock: vi.fn(),
@@ -22,12 +22,61 @@ vi.mock("node:fs", async (importOriginal) => {
2222
const mockGet = vi.mocked(httpsGet);
2323
const tempDirs: string[] = [];
2424

25+
function httpsResponse(
26+
statusCode: number,
27+
headers: Record<string, string> = {},
28+
body?: string | ((url: string) => string),
29+
): typeof httpsGet {
30+
return ((_url: string, callback: (response: IncomingMessage) => void) => {
31+
const response = new PassThrough() as PassThrough & {
32+
statusCode: number;
33+
headers: Record<string, string>;
34+
};
35+
response.statusCode = statusCode;
36+
response.headers = headers;
37+
const request = new EventEmitter() as ClientRequest;
38+
request.setTimeout = vi.fn();
39+
callback(response as unknown as IncomingMessage);
40+
if (body !== undefined) {
41+
queueMicrotask(() => response.end(typeof body === "function" ? body(_url) : body));
42+
}
43+
return request;
44+
}) as typeof httpsGet;
45+
}
46+
2547
afterEach(() => {
2648
vi.clearAllMocks();
2749
for (const dir of tempDirs.splice(0)) rmSync(dir, { recursive: true, force: true });
2850
});
2951

3052
describe("downloadFile", () => {
53+
it("keeps concurrent partial downloads separate", async () => {
54+
mockGet.mockImplementation(httpsResponse(200, {}, (url) => url));
55+
56+
const dir = mkdtempSync(join(tmpdir(), "hyperframes-download-"));
57+
tempDirs.push(dir);
58+
const dest = join(dir, "model.onnx");
59+
const first = "https://example.test/first";
60+
const second = "https://example.test/second";
61+
62+
await Promise.all([downloadFile(first, dest), downloadFile(second, dest)]);
63+
expect([first, second]).toContain(readFileSync(dest, "utf-8"));
64+
});
65+
66+
it("rejects a response that exceeds its byte limit", async () => {
67+
mockGet.mockImplementation(httpsResponse(200, {}, "too large"));
68+
69+
const dir = mkdtempSync(join(tmpdir(), "hyperframes-download-"));
70+
tempDirs.push(dir);
71+
const dest = join(dir, "model.onnx");
72+
73+
await expect(
74+
downloadFile("https://example.test/model.onnx", dest, { maxBytes: 3 }),
75+
).rejects.toThrow("Download exceeded 3 bytes");
76+
expect(existsSync(dest)).toBe(false);
77+
expect(readdirSync(dir)).toEqual([]);
78+
});
79+
3180
it("rejects an idle response and removes the partial file", async () => {
3281
const actualFs = await vi.importActual<typeof import("node:fs")>("node:fs");
3382
let responseClosed = false;
@@ -87,3 +136,59 @@ describe("downloadFile", () => {
87136
expect(existsSync(dest)).toBe(false);
88137
});
89138
});
139+
140+
describe("redirect handling", () => {
141+
it("rejects a redirect that the HTTPS client cannot follow", async () => {
142+
mockGet
143+
.mockImplementationOnce(
144+
httpsResponse(307, { location: "http://cdn.example.test/model.onnx" }),
145+
)
146+
.mockImplementationOnce(() => {
147+
throw new TypeError('Protocol "http:" not supported');
148+
});
149+
150+
const dir = mkdtempSync(join(tmpdir(), "hyperframes-download-"));
151+
tempDirs.push(dir);
152+
await expect(
153+
downloadFile("https://example.test/model.onnx", join(dir, "model.onnx")),
154+
).rejects.toThrow('Protocol "http:" not supported');
155+
});
156+
157+
it("follows every redirect a host may answer with", () => {
158+
// 307 is the one that broke: HuggingFace answers the tokenizer with it,
159+
// and the original set stopped at 302, so the download fell through to the
160+
// not-200 branch and hung.
161+
for (const code of [301, 302, 303, 307, 308]) {
162+
expect(REDIRECT_CODES.has(code)).toBe(true);
163+
}
164+
});
165+
166+
it("does not treat a success or an error as a redirect", () => {
167+
for (const code of [200, 204, 404, 500]) {
168+
expect(REDIRECT_CODES.has(code)).toBe(false);
169+
}
170+
});
171+
172+
it("resolves a relative location against the url that sent it", () => {
173+
// The actual failure: handing this path back as a request target is not a
174+
// valid URL, so nothing was ever fetched.
175+
expect(
176+
redirectTarget(
177+
"/api/resolve-cache/models/x/tokenizer.json",
178+
"https://huggingface.co/x/resolve/main/tokenizer.json",
179+
),
180+
).toBe("https://huggingface.co/api/resolve-cache/models/x/tokenizer.json");
181+
});
182+
183+
it("keeps an absolute location as it is", () => {
184+
expect(redirectTarget("https://cdn.example.com/blob", "https://huggingface.co/a/b")).toBe(
185+
"https://cdn.example.com/blob",
186+
);
187+
});
188+
189+
it("resolves a sibling-relative location", () => {
190+
expect(redirectTarget("other.json", "https://example.com/a/b/file.json")).toBe(
191+
"https://example.com/a/b/other.json",
192+
);
193+
});
194+
});

‎packages/cli/src/utils/download.ts‎

Lines changed: 96 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,37 @@
1+
import { randomUUID } from "node:crypto";
12
import { createWriteStream, renameSync, unlinkSync } from "node:fs";
23
import { get as httpsGet } from "node:https";
3-
import type { IncomingMessage } from "node:http";
4+
import type { ClientRequest, IncomingMessage } from "node:http";
5+
import { Transform } from "node:stream";
46
import { pipeline } from "node:stream/promises";
57

68
const DEFAULT_DOWNLOAD_TIMEOUT_MS = 30_000;
79

810
export interface DownloadOptions {
911
/** Abort after this many milliseconds without network activity. */
1012
timeoutMs?: number;
13+
/** Reject before writing more than this many response bytes. */
14+
maxBytes?: number;
1115
}
1216

17+
/** Every redirect a host may reasonably answer with, not just the two we saw first. */
18+
export const REDIRECT_CODES = new Set([301, 302, 303, 307, 308]);
19+
20+
/**
21+
* Where a redirect actually points.
22+
*
23+
* Split out because this is the part that was wrong, and it is the part that
24+
* can be checked without a socket: hosts answer with relative locations
25+
* (`/api/resolve-cache/...`) far more often than the original code assumed, and
26+
* handing that string back as a request target fails.
27+
*/
28+
export function redirectTarget(location: string, from: string): string {
29+
return new URL(location, from).toString();
30+
}
31+
32+
/** Enough hops for a CDN handoff, few enough that a redirect loop still ends. */
33+
const MAX_REDIRECTS = 10;
34+
1335
function removePartialFile(path: string): void {
1436
try {
1537
unlinkSync(path);
@@ -18,51 +40,97 @@ function removePartialFile(path: string): void {
1840
}
1941
}
2042

43+
function enforceByteLimit(maxBytes: number): Transform {
44+
let received = 0;
45+
return new Transform({
46+
transform(chunk: Buffer, _encoding, callback) {
47+
received += chunk.byteLength;
48+
callback(
49+
received > maxBytes ? new Error(`Download exceeded ${maxBytes} bytes`) : null,
50+
chunk,
51+
);
52+
},
53+
});
54+
}
55+
2156
/**
2257
* Download a file from a URL, following redirects.
2358
* Uses atomic write (download to .tmp, rename on success) to prevent
2459
* corrupt partial files from persisting in the cache on interruption.
60+
*
61+
* Location headers are resolved against the URL that sent them, because they
62+
* are frequently relative: a host answering 307 with `/api/resolve-cache/...`
63+
* is normal, and passing that string back as a request target is not.
64+
*
65+
* The timeout is per request, so it re-arms on each hop of a redirect chain
66+
* rather than budgeting the whole chain. A stalled socket is what it exists to
67+
* catch, and without it the CLI waits forever.
2568
*/
2669
export function downloadFile(
2770
url: string,
2871
dest: string,
2972
options: DownloadOptions = {},
3073
): Promise<void> {
31-
const tmp = `${dest}.tmp`;
74+
const tmp = `${dest}.${process.pid}.${randomUUID()}.tmp`;
3275
const timeoutMs = options.timeoutMs ?? DEFAULT_DOWNLOAD_TIMEOUT_MS;
76+
const maxBytes = options.maxBytes;
3377
return new Promise((resolve, reject) => {
34-
const follow = (u: string) => {
78+
const follow = (u: string, hops = 0) => {
3579
let activeResponse: IncomingMessage | undefined;
3680
let responsePipelineStarted = false;
3781
let requestError: Error | undefined;
38-
const request = httpsGet(u, (res) => {
39-
activeResponse = res;
40-
if (res.statusCode === 301 || res.statusCode === 302) {
41-
const location = res.headers.location;
42-
if (location) {
82+
let request: ClientRequest;
83+
try {
84+
request = httpsGet(u, (res) => {
85+
activeResponse = res;
86+
if (res.statusCode && REDIRECT_CODES.has(res.statusCode)) {
87+
const location = res.headers.location;
88+
if (location) {
89+
if (hops >= MAX_REDIRECTS) {
90+
res.resume();
91+
removePartialFile(tmp);
92+
reject(
93+
new Error(`Download failed: more than ${MAX_REDIRECTS} redirects from ${url}`),
94+
);
95+
return;
96+
}
97+
res.resume();
98+
try {
99+
follow(redirectTarget(location, u), hops + 1);
100+
} catch (error) {
101+
removePartialFile(tmp);
102+
reject(error);
103+
}
104+
return;
105+
}
106+
}
107+
if (res.statusCode !== 200) {
43108
res.resume();
44-
follow(location);
109+
removePartialFile(tmp);
110+
reject(new Error(`Download failed: HTTP ${res.statusCode}`));
45111
return;
46112
}
47-
}
48-
if (res.statusCode !== 200) {
49-
res.resume();
50-
removePartialFile(tmp);
51-
reject(new Error(`Download failed: HTTP ${res.statusCode}`));
52-
return;
53-
}
54-
const file = createWriteStream(tmp);
55-
responsePipelineStarted = true;
56-
pipeline(res, file)
57-
.then(() => {
58-
renameSync(tmp, dest);
59-
resolve();
60-
})
61-
.catch((err) => {
62-
removePartialFile(tmp);
63-
reject(requestError ?? err);
64-
});
65-
});
113+
const file = createWriteStream(tmp);
114+
responsePipelineStarted = true;
115+
const transfer =
116+
maxBytes === undefined
117+
? pipeline(res, file)
118+
: pipeline(res, enforceByteLimit(maxBytes), file);
119+
transfer
120+
.then(() => {
121+
renameSync(tmp, dest);
122+
resolve();
123+
})
124+
.catch((err) => {
125+
removePartialFile(tmp);
126+
reject(requestError ?? err);
127+
});
128+
});
129+
} catch (error) {
130+
removePartialFile(tmp);
131+
reject(error);
132+
return;
133+
}
66134
request.setTimeout(timeoutMs, () => {
67135
request.destroy(new Error(`Download timed out after ${timeoutMs}ms`));
68136
});

0 commit comments

Comments
 (0)