quak/src/download/index.ts

import { randomUUID } from "node:crypto";
import { rename, rm, writeFile } from "node:fs/promises";
import { dirname, join } from "node:path";
import {
    fromBase64,
    initStreamPull,
    pullStreamChunk,
    STREAM_CHUNK_OVERHEAD,
    STREAM_CHUNK_SIZE,
    streamTagFinal,
} from "../crypto/index.js";
import { TruncatedStreamError } from "../errors.js";
import { withRetry } from "../retry.js";
import type { ApiClient } from "../api/client.js";
import type { EnteFile } from "../model/types.js";

export interface DownloadResult {
    path: string;
    bytesWritten: number;
}

const ENC_CHUNK_SIZE = STREAM_CHUNK_SIZE + STREAM_CHUNK_OVERHEAD;

const streamDecrypt = async (
    stream: ReadableStream<Uint8Array>,
    header: Uint8Array,
    key: Uint8Array,
): Promise<Uint8Array> => {
    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;

    for (;;) {
        const { done, value } = await reader.read();
        if (value) {
            const merged = new Uint8Array(buffer.length + value.length);
            merged.set(buffer);
            merged.set(value, buffer.length);
            buffer = merged;
        }

        while (buffer.length >= ENC_CHUNK_SIZE) {
            const encChunk = buffer.slice(0, ENC_CHUNK_SIZE);
            buffer = buffer.slice(ENC_CHUNK_SIZE);
            const { plaintext, tag } = pullStreamChunk(state, encChunk);
            plainChunks.push(plaintext);
            totalPlain += plaintext.length;
            chunksPulled++;
            lastTag = tag;
        }

        if (done) {
            if (buffer.length > 0) {
                // Whatever is left over once every whole chunk has been
                // consumed must be the stream's final chunk, and a final
                // chunk that actually arrived in full authenticates. If it
                // does not, the body stopped part-way through a chunk — the
                // 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.
                let pulled;
                try {
                    pulled = pullStreamChunk(state, buffer);
                } catch (err) {
                    throw new TruncatedStreamError(
                        `download: stream truncated: response body ended with ${buffer.length} trailing bytes that did not authenticate as a final chunk (transfer stopped mid-chunk, or the data is corrupt)`,
                        { cause: err },
                    );
                }
                plainChunks.push(pulled.plaintext);
                totalPlain += pulled.plaintext.length;
                chunksPulled++;
                lastTag = pulled.tag;
            }
            break;
        }
    }

    // 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",
        );
    }
    const tagFinal = streamTagFinal();
    if (lastTag !== tagFinal) {
        throw new TruncatedStreamError(
            `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;
};

// Write `plaintext` to `destination` atomically: 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.
const writeAtomic = async (
    destination: string,
    plaintext: Uint8Array,
): Promise<void> => {
    // The random suffix keeps concurrent downloads of the same destination
    // from stepping on each other's temporary file.
    const tmpPath = join(dirname(destination), `.quak-${randomUUID()}.tmp`);
    try {
        await writeFile(tmpPath, plaintext);
        await rename(tmpPath, destination);
    } catch (err) {
        // Best-effort cleanup. A failure to remove the temporary file must
        // never replace the error that actually explains what went wrong.
        await rm(tmpPath, { force: true }).catch(() => undefined);
        throw err;
    }
};

// Fetch a stream and decrypt it, 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
// mid-body — the dominant failure mode for multi-megabyte photos over a CDN —
// throws in `streamDecrypt` and never reaches `ApiClient` at all. Retrying the
// request alone would miss it entirely.
//
// The client's own retry is therefore switched off for these two calls: with
// both layers active the budgets would multiply, and the library default of
// 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.
const fetchAndDecrypt = async (
    api: ApiClient,
    openStream: () => Promise<ReadableStream<Uint8Array>>,
    header: Uint8Array,
    key: Uint8Array,
): Promise<Uint8Array> =>
    withRetry(async () => {
        const stream = await openStream();
        return streamDecrypt(stream, header, key);
    }, api.getRetryOptions());

export const downloadFile = async (
    api: ApiClient,
    file: EnteFile,
    outPath?: string,
): Promise<DownloadResult> => {
    const resolvedPath = outPath ?? file.metadata.title;
    const header = fromBase64(file.file.decryptionHeader);
    const plaintext = await fetchAndDecrypt(
        api,
        () => api.getFileStream(file.id, { retry: false }),
        header,
        file.key,
    );
    // 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 };
};

export const downloadThumbnail = async (
    api: ApiClient,
    file: EnteFile,
    outPath?: string,
): Promise<DownloadResult> => {
    const resolvedPath = outPath ?? `thumb_${file.metadata.title}`;
    const header = fromBase64(file.thumbnail.decryptionHeader);
    const plaintext = await fetchAndDecrypt(
        api,
        () => api.getThumbnailStream(file.id, { retry: false }),
        header,
        file.key,
    );
    await writeAtomic(resolvedPath, plaintext);
    return { path: resolvedPath, bytesWritten: plaintext.length };
};