sneak/quak
1
0
Fork 0
Code Issues 8 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Stream decrypted downloads to disk with bounded memory (closes #40)
check / check (push) Successful in 32s
Details

Browse Source 
Originals no longer buffer the whole decrypted file in RAM: streamDecrypt hands each secretstream chunk to a sink and the download path writes it to the staged temp file, so peak memory is one chunk regardless of file size. The atomic write moved inside the retry loop; only a TAG_FINAL-authenticated attempt renames; truncation leaves no destination file. Does not close #21 (the streamDecrypt accumulation-buffer recopy stays open).

Model: opus-4-8
This commit was merged in pull request #59.
This commit is contained in:
clawbot
2026-09-22 13:54:20 +02:00
parent 72ea8dcb01
commit 4b4f550f89
2 changed files with 234 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/download/index.ts
+102 -47
View File
@@ -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<Uint8Array>,
header: Uint8Array,
key: Uint8Array,
sink: (plaintext: Uint8Array) => Promise<void>,
onProgress?: ProgressCallback,
): Promise<Uint8Array> => {
): Promise<number> => {
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<void> => {
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<void>,
): Promise<void> => {
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<void> =>
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<Uint8Array>,
header: Uint8Array,
key: Uint8Array,
onProgress?: ProgressCallback,
): Promise<number> => {
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<ReadableStream<Uint8Array>>,
header: Uint8Array,
key: Uint8Array,
destination: string,
onProgress?: ProgressCallback,
): Promise<Uint8Array> =>
): Promise<number> =>
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<DownloadResult> => {
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<DownloadResult> => {
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 };
};
test/download/download.test.ts
+132 -4
View File
@@ -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<typeof import("node:fs/promises")>();
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<unknown> => {
if (data instanceof Uint8Array) {
writeHook.writes.push({
path: String(path),
length: data.length,
});
}
return (realWrite as (...a: unknown[]) => Promise<unknown>)(
data,
...rest2,
);
}) as typeof handle.write;
return handle;
},
rename: async (from: string, to: string): Promise<void> => {
@@ -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