diff --git a/src/download/index.ts b/src/download/index.ts index d0be2bd..fa53b21 100644 --- a/src/download/index.ts +++ b/src/download/index.ts @@ -1,5 +1,6 @@ import { randomUUID } from "node:crypto"; import { open, rename, rm } from "node:fs/promises"; +import type { FileHandle } from "node:fs/promises"; import { dirname, join } from "node:path"; import { fromBase64, @@ -28,20 +29,44 @@ export type ProgressCallback = (bytesDone: number) => void; const ENC_CHUNK_SIZE = STREAM_CHUNK_SIZE + STREAM_CHUNK_OVERHEAD; +// Decrypt a secretstream body, handing each plaintext chunk to `sink` as it is +// produced rather than accumulating the whole file. Peak memory is one +// ciphertext chunk of network buffer plus one plaintext chunk — bounded by +// `STREAM_CHUNK_SIZE` regardless of the file's size — so a multi-gigabyte video +// no longer needs its size again in RAM. Returns the total plaintext length. +// +// The truncation contract is exactly the buffered version's, only the sink is +// new: a body cut short still decrypts and authenticates up to its last whole +// chunk, so the absence of TAG_FINAL is the sole evidence it was cut short, and +// this throws rather than let a caller keep a short file. The sink has already +// seen those chunks by then; the caller (`decryptToTemp`) stages them in a temp +// file that is renamed into place only on a clean return, so a throw leaves +// nothing on disk. const streamDecrypt = async ( stream: ReadableStream, header: Uint8Array, key: Uint8Array, + sink: (plaintext: Uint8Array) => Promise, onProgress?: ProgressCallback, -): Promise => { +): Promise => { const state = initStreamPull(header, key); const reader = stream.getReader(); let buffer = new Uint8Array(0); - const plainChunks: Uint8Array[] = []; let totalPlain = 0; let chunksPulled = 0; let lastTag = -1; + const consume = async ( + plaintext: Uint8Array, + tag: number, + ): Promise => { + await sink(plaintext); + totalPlain += plaintext.length; + chunksPulled++; + lastTag = tag; + onProgress?.(totalPlain); + }; + for (;;) { const { done, value } = await reader.read(); if (value) { @@ -54,12 +79,10 @@ const streamDecrypt = async ( while (buffer.length >= ENC_CHUNK_SIZE) { const encChunk = buffer.slice(0, ENC_CHUNK_SIZE); buffer = buffer.slice(ENC_CHUNK_SIZE); + // A whole chunk that fails to authenticate while the stream carries + // on is corruption, not truncation; that error propagates unchanged. const { plaintext, tag } = pullStreamChunk(state, encChunk); - plainChunks.push(plaintext); - totalPlain += plaintext.length; - chunksPulled++; - lastTag = tag; - onProgress?.(totalPlain); + await consume(plaintext, tag); } if (done) { @@ -71,7 +94,9 @@ const streamDecrypt = async ( // ordinary shape of a dropped connection. Poly1305 cannot // tell a partial chunk from a corrupt one, so this is // reported as the truncation it almost always is, with the - // authentication failure kept as the error's cause. + // authentication failure kept as the error's cause. Only the + // pull is guarded: a sink failure on a chunk that did + // authenticate is a disk error, not a truncation. let pulled; try { pulled = pullStreamChunk(state, buffer); @@ -81,11 +106,7 @@ const streamDecrypt = async ( { cause: err }, ); } - plainChunks.push(pulled.plaintext); - totalPlain += pulled.plaintext.length; - chunksPulled++; - lastTag = pulled.tag; - onProgress?.(totalPlain); + await consume(pulled.plaintext, pulled.tag); } break; } @@ -94,8 +115,6 @@ const streamDecrypt = async ( // Only the last chunk of a secretstream carries TAG_FINAL. Everything a // dropped connection did deliver still decrypts and authenticates, so the // absence of TAG_FINAL is the only evidence that the body was cut short. - // Returning a short plaintext here would put a corrupt file on disk that - // later backup runs would treat as complete. if (chunksPulled === 0) { throw new TruncatedStreamError( "download: stream truncated: response body contained no secretstream chunks", @@ -107,21 +126,16 @@ const streamDecrypt = async ( `download: stream truncated: last chunk tag ${lastTag}, expected TAG_FINAL (${tagFinal})`, ); } - - const result = new Uint8Array(totalPlain); - let offset = 0; - for (const chunk of plainChunks) { - result.set(chunk, offset); - offset += chunk.length; - } - return result; + return totalPlain; }; -// Write `plaintext` to `destination` atomically and durably: stage it in a -// temporary sibling file (same directory, so the rename cannot cross a -// filesystem boundary) and rename it into place. Callers therefore never -// observe a partially written destination, and a pre-existing file at that -// path is replaced only once the new contents are complete on disk. +// Stage a write to `destination` atomically and durably, then rename it into +// place. `fill` writes the contents into the open temp file handle — either the +// whole buffer at once (`writeAtomic`) or chunk by chunk as they decrypt +// (`decryptToTemp`). The temp file is a sibling of the destination (same +// directory, so the rename cannot cross a filesystem boundary), so callers +// never observe a partially written destination, and a pre-existing file is +// replaced only once the new contents are complete on disk. // // Durability against a power cut needs two fsyncs. Without them the write can // return while the data or the rename is still only in the kernel's page @@ -131,10 +145,12 @@ const streamDecrypt = async ( // is fsynced after it, so both the bytes and the new directory entry are on // stable storage before this returns. // -// Exported so the metadata store can reuse the same durable write. -export const writeAtomic = async ( +// On any failure — including a `fill` that throws because the stream was +// truncated — the temp file is removed, so the destination is untouched and no +// scratch file is left to fill the disk on repeated failures. +const stageAtomic = async ( destination: string, - plaintext: Uint8Array, + fill: (handle: FileHandle) => Promise, ): Promise => { const dir = dirname(destination); // The random suffix keeps concurrent downloads of the same destination @@ -143,7 +159,7 @@ export const writeAtomic = async ( try { const handle = await open(tmpPath, "w"); try { - await handle.writeFile(plaintext); + await fill(handle); await handle.sync(); } finally { await handle.close(); @@ -166,7 +182,44 @@ export const writeAtomic = async ( } }; -// Fetch a stream and decrypt it, retrying the whole sequence. +// Write `plaintext` to `destination` atomically and durably. Exported so the +// metadata store can reuse the same durable write for small whole-buffer +// payloads; originals go through `decryptToTemp` instead so they never buffer. +export const writeAtomic = async ( + destination: string, + plaintext: Uint8Array, +): Promise => + stageAtomic(destination, (handle) => handle.writeFile(plaintext)); + +// Decrypt `stream` straight to `destination`, one plaintext chunk at a time, +// under the atomic writer's temp-then-rename discipline. Memory stays bounded +// by the chunk size: each decrypted chunk is written to the temp file and +// dropped. The rename happens only after the stream authenticates as terminated +// on TAG_FINAL; a truncated stream throws and leaves the destination untouched. +// Returns the plaintext length written. +const decryptToTemp = async ( + destination: string, + stream: ReadableStream, + header: Uint8Array, + key: Uint8Array, + onProgress?: ProgressCallback, +): Promise => { + let bytesWritten = 0; + await stageAtomic(destination, async (handle) => { + bytesWritten = await streamDecrypt( + stream, + header, + key, + async (plaintext) => { + await handle.write(plaintext); + }, + onProgress, + ); + }); + return bytesWritten; +}; + +// Fetch a stream and decrypt it to `destination`, retrying the whole sequence. // // The request is only the first third of a download. `getXStream` returns as // soon as headers arrive, and the bytes are pulled here, so a socket reset @@ -179,18 +232,24 @@ export const writeAtomic = async ( // four attempts would mean sixteen requests for one file. The policy comes // from the client so a caller that configured one gets it here too. // -// A retry starts the file over from byte zero: the secretstream pull state is -// not resumable and there is no Range support on these endpoints. +// Because the plaintext is streamed to disk rather than buffered, the atomic +// write is part of the retried unit. A retry starts the file over from byte +// zero — the secretstream pull state is not resumable and there is no Range +// support — staging into a fresh temp file each time: a failed attempt writes +// and then removes its own temp file, and only the attempt that reaches +// TAG_FINAL renames one into place, so a download that needed three tries still +// performs exactly one rename over the destination. const fetchAndDecrypt = async ( api: ApiClient, openStream: () => Promise>, header: Uint8Array, key: Uint8Array, + destination: string, onProgress?: ProgressCallback, -): Promise => +): Promise => withRetry(async () => { const stream = await openStream(); - return streamDecrypt(stream, header, key, onProgress); + return decryptToTemp(destination, stream, header, key, onProgress); }, api.getRetryOptions()); export const downloadFile = async ( @@ -201,19 +260,15 @@ export const downloadFile = async ( ): Promise => { const resolvedPath = outPath ?? file.metadata.title; const header = fromBase64(file.file.decryptionHeader); - const plaintext = await fetchAndDecrypt( + const bytesWritten = await fetchAndDecrypt( api, () => api.getFileStream(file.id, { retry: false }), header, file.key, + resolvedPath, onProgress, ); - // Outside the retry, deliberately: only the attempt that produced a - // complete, authenticated plaintext gets to stage a temporary file, so a - // download that needed three tries still performs exactly one write and - // one rename. - await writeAtomic(resolvedPath, plaintext); - return { path: resolvedPath, bytesWritten: plaintext.length }; + return { path: resolvedPath, bytesWritten }; }; export const downloadThumbnail = async ( @@ -224,13 +279,13 @@ export const downloadThumbnail = async ( ): Promise => { const resolvedPath = outPath ?? `thumb_${file.metadata.title}`; const header = fromBase64(file.thumbnail.decryptionHeader); - const plaintext = await fetchAndDecrypt( + const bytesWritten = await fetchAndDecrypt( api, () => api.getThumbnailStream(file.id, { retry: false }), header, file.key, + resolvedPath, onProgress, ); - await writeAtomic(resolvedPath, plaintext); - return { path: resolvedPath, bytesWritten: plaintext.length }; + return { path: resolvedPath, bytesWritten }; }; diff --git a/test/download/download.test.ts b/test/download/download.test.ts index 95bcdf2..6eea19d 100644 --- a/test/download/download.test.ts +++ b/test/download/download.test.ts @@ -118,6 +118,21 @@ const durabilityHook = vi.hoisted(() => ({ events: [] as string[], })); +/** + * `FileHandle.write` is wrapped so the tests can watch the streaming decrypt + * path put plaintext on disk one chunk at a time. This is the direct evidence + * that memory is bounded by the chunk size and not the file size: a buffered + * downloader would hand the whole file to a single write, whereas the streaming + * one issues one write per secretstream chunk, none larger than + * `STREAM_CHUNK_SIZE`. Each write records the temp path it targeted and its + * length. `writeFile` (which the whole-buffer `writeAtomic` uses) is a distinct + * native call and does not go through this method, so only the streaming path + * is observed here. + */ +const writeHook = vi.hoisted(() => ({ + writes: [] as { path: string; length: number }[], +})); + vi.mock("node:fs/promises", async (importOriginal) => { const actual = await importOriginal(); const { existsSync: sourceExists } = await import("node:fs"); @@ -138,6 +153,22 @@ vi.mock("node:fs/promises", async (importOriginal) => { durabilityHook.events.push(`sync:${String(flags)}:${path}`); await realSync(); }; + const realWrite = handle.write.bind(handle); + handle.write = (async ( + data: unknown, + ...rest2: unknown[] + ): Promise => { + if (data instanceof Uint8Array) { + writeHook.writes.push({ + path: String(path), + length: data.length, + }); + } + return (realWrite as (...a: unknown[]) => Promise)( + data, + ...rest2, + ); + }) as typeof handle.write; return handle; }, rename: async (from: string, to: string): Promise => { @@ -159,6 +190,7 @@ beforeEach(() => { renameHook.calls.length = 0; renameHook.failWith = null; durabilityHook.events.length = 0; + writeHook.writes.length = 0; }); let testDir: string; @@ -947,10 +979,13 @@ describe.each(entryPoints)("$name retries", ({ name, download }) => { }); it("stages one temp file for the attempt that succeeded, not one per attempt", async () => { - // The atomic write stays outside the retry loop. A retried download - // must not leave a trail of half-written scratch files, and the - // destination must be touched exactly once — by the attempt that - // produced a complete, authenticated plaintext. + // Each streaming attempt stages into its own temp file, but a retried + // download must not leave a trail of half-written scratch files: a + // failed attempt removes its temp file, and the destination is renamed + // into place exactly once — by the attempt that produced a complete, + // authenticated plaintext. (Here the two failed attempts reset before a + // whole chunk is pulled, so they write nothing; the point stands either + // way — see the retry-restart test below, where they do write.) const { key, header, ciphertext } = smallFixture(42); const { fetch } = scriptedCdnFetch( { kind: "reset", bytes: ciphertext.slice(0, 16) }, @@ -1064,6 +1099,99 @@ describe.each(entryPoints)("$name retries", ({ name, download }) => { }); }); +// --------------------------------------------------------------------------- +// Streaming decrypt to disk +// +// The plaintext is never held whole in memory: each secretstream chunk is +// written to the temp file as it is decrypted, so peak memory is bounded by the +// chunk size rather than the file size. These tests watch the writes directly +// (see `writeHook`) rather than infer memory behaviour from the final file. +// --------------------------------------------------------------------------- + +describe.each(entryPoints)("$name streams to disk", ({ name, download }) => { + const freshDir = (): string => + mkdtempSync(join(testDir, `${name}-stream-`)); + + /** Writes recorded against staged temp files (not the `writeFile` path). */ + const tempWrites = (): { path: string; length: number }[] => + writeHook.writes.filter((w) => w.path.endsWith(".tmp")); + + it("writes one chunk at a time, none larger than STREAM_CHUNK_SIZE", async () => { + // The multi-chunk fixture decrypts to one full 4 MiB chunk plus a small + // final chunk. A streaming writer therefore issues exactly two writes, + // of STREAM_CHUNK_SIZE and then the final chunk's length — never a + // single write carrying the whole 4 MiB + 1 KiB file. That per-chunk + // shape is what "memory bounded by chunk size" means in practice: the + // plaintext is handed to the filesystem and dropped, chunk by chunk. + const { api, file } = fixtureFor( + multiChunkKey, + multiChunk.header, + multiChunk.body, + ); + const outPath = join(freshDir(), "streamed.bin"); + + const result = await download(api, file, outPath); + + const writes = tempWrites(); + expect(writes.map((w) => w.length)).toEqual([ + STREAM_CHUNK_SIZE, + multiChunk.plaintext.length - STREAM_CHUNK_SIZE, + ]); + // No single write ever carried the whole file, and every write fits in + // one chunk's worth of memory. + for (const w of writes) { + expect(w.length).toBeLessThanOrEqual(STREAM_CHUNK_SIZE); + } + expect(result.bytesWritten).toBe(multiChunk.plaintext.length); + expectSameBytes(readFileSync(outPath), multiChunk.plaintext); + }); + + it("restarts from byte zero on a retry, replacing the temp file cleanly", async () => { + // The secretstream pull state is not resumable, so a retry cannot + // continue a half-written file — it must start over. The first attempt + // here delivers a complete leading chunk and then stops before the + // TAG_FINAL chunk: 4 MiB of plaintext lands in a temp file, then the + // download is rejected as truncated and that temp file is discarded. + // The retry streams the whole body into a *fresh* temp file, so the + // destination ends up with exactly the plaintext once — never the + // leading chunk twice, and never a stale temp file left behind. + const truncatedBody = multiChunk.body.slice( + 0, + multiChunk.finalChunkOffset, + ); + const { fetch, requests } = scriptedCdnFetch( + { kind: "body", bytes: truncatedBody }, + { kind: "body", bytes: multiChunk.body }, + ); + const api = new ApiClient({ + fetch, + retry: { ...noWait, attempts: 4 }, + }); + const file = buildMockEnteFile( + multiChunkKey, + multiChunk.header, + multiChunk.header, + ); + const dir = freshDir(); + const outPath = join(dir, "retry-restart.bin"); + + const result = await download(api, file, outPath); + + expect(requests()).toBe(2); + // Both attempts streamed to disk, each into its own temp file: the + // truncated first attempt wrote before it failed, proving the retry did + // not resume a partial file but replaced it. + const distinctTemps = new Set(tempWrites().map((w) => w.path)); + expect(distinctTemps.size).toBe(2); + // The destination holds the complete plaintext exactly once, and no + // temp file survives. + expect(result.bytesWritten).toBe(multiChunk.plaintext.length); + expectSameBytes(readFileSync(outPath), multiChunk.plaintext); + expect(renameHook.calls).toHaveLength(1); + expect(readdirSync(dir)).toEqual(["retry-restart.bin"]); + }); +}); + describe("download retries: corruption is not retried", () => { it("gives up immediately on a chunk that failed to authenticate", async () => { // A whole chunk that failed to authenticate while the stream