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,
    STREAM_TAG_FINAL,
} from "../crypto/index.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) {
                const { plaintext, tag } = pullStreamChunk(state, buffer);
                plainChunks.push(plaintext);
                totalPlain += plaintext.length;
                chunksPulled++;
                lastTag = 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 Error(
            "download: stream truncated: response body contained no secretstream chunks",
        );
    }
    if (lastTag !== STREAM_TAG_FINAL) {
        throw new Error(
            `download: stream truncated: last chunk tag ${lastTag}, expected TAG_FINAL (${STREAM_TAG_FINAL})`,
        );
    }

    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;
    }
};

export const downloadFile = async (
    api: ApiClient,
    file: EnteFile,
    outPath?: string,
): Promise<DownloadResult> => {
    const resolvedPath = outPath ?? file.metadata.title;
    const stream = await api.getFileStream(file.id);
    const header = fromBase64(file.file.decryptionHeader);
    const plaintext = await streamDecrypt(stream, header, file.key);
    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 stream = await api.getThumbnailStream(file.id);
    const header = fromBase64(file.thumbnail.decryptionHeader);
    const plaintext = await streamDecrypt(stream, header, file.key);
    await writeAtomic(resolvedPath, plaintext);
    return { path: resolvedPath, bytesWritten: plaintext.length };
};