quak/test/download/download.test.ts

/**
 * Tests for `downloadFile` and `downloadThumbnail`.
 *
 * These functions combine the ApiClient's streaming download with the
 * secretstream pull decryption to recover the plaintext file content and
 * write it to disk.
 *
 * The encrypted body returned by the CDN is a concatenation of
 * secretstream ciphertext chunks. Each chunk is at most
 * `STREAM_CHUNK_SIZE + STREAM_CHUNK_OVERHEAD` bytes (4 MiB + 17 bytes).
 * The download function buffers incoming network data, splits it on the
 * chunk boundary, and feeds each piece to `pullStreamChunk`.
 *
 * Two contracts are load-bearing for anyone using this library as a backup
 * tool, and both are documented by the tests below:
 *
 * 1. **Truncation is an error, never a short file.** Only the final chunk of
 *    a secretstream carries `TAG_FINAL`. A download cut short by a dropped
 *    connection still decrypts cleanly up to the last whole chunk, so without
 *    an explicit `TAG_FINAL` check a truncated body is indistinguishable from
 *    a complete one. `streamDecrypt` therefore refuses to return unless the
 *    stream ended on `TAG_FINAL`, and the error says the stream was truncated.
 *
 * 2. **The destination path is written atomically.** Plaintext goes to a
 *    temporary sibling file first and is `rename`d into place only after the
 *    whole stream has decrypted and verified. A caller that sees no exception
 *    can rely on the destination containing the complete, authenticated file;
 *    a caller that sees an exception can rely on the destination being
 *    untouched — whatever was there before is still there, byte for byte, and
 *    no partial file has appeared. This matters because `runBackup` skips any
 *    existing non-empty file, so a partial write would be treated as complete
 *    forever after.
 *
 * These tests build synthetic encrypted files using sodium's push API,
 * serve them from a mock fetch, and verify the decrypted output on disk.
 */

import {
    existsSync,
    readdirSync,
    readFileSync,
    rmSync,
    mkdtempSync,
    writeFileSync,
} from "node:fs";
import { join } from "node:path";
import { tmpdir } from "node:os";
import { createHash } from "node:crypto";
import sodium from "libsodium-wrappers-sumo";
import { beforeAll, afterAll, describe, expect, it } from "vitest";
import { init, toBase64, STREAM_CHUNK_SIZE } from "../../src/crypto/index.js";
import { ApiClient } from "../../src/api/client.js";
import { downloadFile, downloadThumbnail } from "../../src/download/index.js";
import type { EnteFile, FileMetadata } from "../../src/model/types.js";

// ---------------------------------------------------------------------------
// Test helpers
// ---------------------------------------------------------------------------

let testDir: string;

beforeAll(async () => {
    await init();
    await sodium.ready;
    testDir = mkdtempSync(join(tmpdir(), "quak-test-"));
});

afterAll(() => {
    if (testDir && existsSync(testDir)) {
        rmSync(testDir, { recursive: true, force: true });
    }
});

/**
 * Encrypt `plaintext` as a secretstream file body (single chunk with
 * TAG_FINAL). Returns the key, header, and ciphertext that the mock CDN
 * will serve.
 */
const encryptFileBody = (
    plaintext: Uint8Array,
    key: Uint8Array,
): { header: Uint8Array; ciphertext: Uint8Array } => {
    const push = sodium.crypto_secretstream_xchacha20poly1305_init_push(key);
    const ciphertext = sodium.crypto_secretstream_xchacha20poly1305_push(
        push.state,
        plaintext,
        null,
        sodium.crypto_secretstream_xchacha20poly1305_TAG_FINAL,
    );
    return { header: push.header, ciphertext };
};

/**
 * Encrypt a body that spans more than one secretstream chunk, the way the
 * server does for files larger than the 4 MiB plaintext chunk size.
 *
 * Framing matters here: the downloader splits the byte stream on fixed
 * `STREAM_CHUNK_SIZE + STREAM_CHUNK_OVERHEAD` boundaries, so every chunk
 * except the last must carry exactly `STREAM_CHUNK_SIZE` plaintext bytes.
 * Only the last chunk is tagged `TAG_FINAL`; the leading ones are
 * `TAG_MESSAGE`.
 *
 * Returns the header, the concatenated body, the plaintext it decrypts to,
 * and `finalChunkOffset` — the byte offset at which the `TAG_FINAL` chunk
 * begins, so a test can slice it off to simulate a connection that dropped
 * before the end of the file.
 */
const encryptMultiChunkBody = (
    key: Uint8Array,
    leadingChunks: number,
    finalChunkPlainSize: number,
): {
    header: Uint8Array;
    body: Uint8Array;
    plaintext: Uint8Array;
    finalChunkOffset: number;
} => {
    const push = sodium.crypto_secretstream_xchacha20poly1305_init_push(key);
    const cipherParts: Uint8Array[] = [];
    const plainParts: Uint8Array[] = [];

    for (let i = 0; i < leadingChunks; i++) {
        const plain = sodium.randombytes_buf(STREAM_CHUNK_SIZE);
        plainParts.push(plain);
        cipherParts.push(
            sodium.crypto_secretstream_xchacha20poly1305_push(
                push.state,
                plain,
                null,
                sodium.crypto_secretstream_xchacha20poly1305_TAG_MESSAGE,
            ),
        );
    }

    const finalPlain = sodium.randombytes_buf(finalChunkPlainSize);
    plainParts.push(finalPlain);
    const finalCipher = sodium.crypto_secretstream_xchacha20poly1305_push(
        push.state,
        finalPlain,
        null,
        sodium.crypto_secretstream_xchacha20poly1305_TAG_FINAL,
    );

    const finalChunkOffset = cipherParts.reduce((n, c) => n + c.length, 0);
    cipherParts.push(finalCipher);

    return {
        header: push.header,
        body: concat(cipherParts),
        plaintext: concat(plainParts),
        finalChunkOffset,
    };
};

const concat = (parts: Uint8Array[]): Uint8Array => {
    const total = parts.reduce((n, p) => n + p.length, 0);
    const out = new Uint8Array(total);
    let offset = 0;
    for (const p of parts) {
        out.set(p, offset);
        offset += p.length;
    }
    return out;
};

const buildMockEnteFile = (
    key: Uint8Array,
    fileHeader: Uint8Array,
    thumbHeader: Uint8Array,
): EnteFile => ({
    id: 999,
    collectionID: 1,
    ownerID: 1,
    key,
    metadata: {
        title: "test-photo.jpg",
        fileType: "image",
        creationTime: 0,
        modificationTime: 0,
    } as FileMetadata,
    file: { decryptionHeader: toBase64(fileHeader) },
    thumbnail: { decryptionHeader: toBase64(thumbHeader) },
    updationTime: 0,
});

const mockFetchForBody = (body: Uint8Array) => {
    const fake = async (): Promise<Response> =>
        new Response(body, { status: 200 });
    return fake as typeof globalThis.fetch;
};

/**
 * Encrypt a body consisting of one chunk that is *not* tagged TAG_FINAL.
 *
 * This is the cheap way to present a stream that ended without its final
 * chunk: the downloader pulls it, authenticates it, and finds the stream
 * over on a TAG_MESSAGE chunk — the same terminal condition as a large file
 * whose last chunk was lost, without paying for a 4 MiB fixture. The
 * multi-chunk fixture above covers the realistic wire shape; this one is
 * used where the test is really about what happens on disk afterwards.
 */
const encryptNonFinalBody = (
    plaintext: Uint8Array,
    key: Uint8Array,
): { header: Uint8Array; ciphertext: Uint8Array } => {
    const push = sodium.crypto_secretstream_xchacha20poly1305_init_push(key);
    const ciphertext = sodium.crypto_secretstream_xchacha20poly1305_push(
        push.state,
        plaintext,
        null,
        sodium.crypto_secretstream_xchacha20poly1305_TAG_MESSAGE,
    );
    return { header: push.header, ciphertext };
};

/**
 * Compare file contents by digest rather than with `toEqual`. Vitest's deep
 * equality walks multi-megabyte buffers byte by byte, which costs seconds on
 * the 4 MiB fixtures; a digest comparison is exact and effectively free.
 */
const expectSameBytes = (actual: Uint8Array, expected: Uint8Array): void => {
    expect(actual.length).toBe(expected.length);
    expect(createHash("sha256").update(actual).digest("hex")).toBe(
        createHash("sha256").update(expected).digest("hex"),
    );
};

/**
 * A multi-chunk fixture shared by the truncation tests. Building it costs a
 * few MiB of encryption, so it is built once: one full 4 MiB `TAG_MESSAGE`
 * chunk followed by a small `TAG_FINAL` chunk.
 */
let multiChunk: ReturnType<typeof encryptMultiChunkBody>;
let multiChunkKey: Uint8Array;

beforeAll(() => {
    multiChunkKey = sodium.crypto_secretstream_xchacha20poly1305_keygen();
    multiChunk = encryptMultiChunkBody(multiChunkKey, 1, 1024);
});

/**
 * Build an EnteFile plus ApiClient whose file *and* thumbnail streams both
 * serve `body` under `header`. The download path under test is otherwise
 * identical for the two, so every truncation/atomicity case below runs
 * against both entry points from a single fixture.
 */
const fixtureFor = (
    key: Uint8Array,
    header: Uint8Array,
    body: Uint8Array,
): { api: ApiClient; file: EnteFile } => ({
    api: new ApiClient({ fetch: mockFetchForBody(body) }),
    file: buildMockEnteFile(key, header, header),
});

// The two entry points share `streamDecrypt` and the atomic-write wrapper,
// so the contract tests are written once and run against both.
const entryPoints = [
    { name: "downloadFile", download: downloadFile },
    { name: "downloadThumbnail", download: downloadThumbnail },
];

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

describe("downloadFile", () => {
    it("downloads, decrypts, and writes a single-chunk file", async () => {
        const plaintext = new TextEncoder().encode(
            "Hello from quak! This is a test photo payload.",
        );
        const key = sodium.crypto_secretstream_xchacha20poly1305_keygen();
        const { header, ciphertext } = encryptFileBody(plaintext, key);

        // Separate header for thumbnail (not used in this test path but
        // needed to construct the EnteFile)
        const thumbPush =
            sodium.crypto_secretstream_xchacha20poly1305_init_push(key);

        const file = buildMockEnteFile(key, header, thumbPush.header);
        const api = new ApiClient({ fetch: mockFetchForBody(ciphertext) });
        const outPath = join(testDir, "single-chunk.jpg");

        const result = await downloadFile(api, file, outPath);

        // The whole DownloadResult shape is asserted, not just its fields:
        // callers depend on `path` being the destination they asked for
        // (never the temporary file used along the way) and on
        // `bytesWritten` being the plaintext length.
        expect(result).toEqual({
            path: outPath,
            bytesWritten: plaintext.length,
        });
        expect(readFileSync(outPath)).toEqual(Buffer.from(plaintext));
    });

    it("uses metadata.title as filename when outPath is omitted", async () => {
        const plaintext = new Uint8Array([1, 2, 3]);
        const key = sodium.crypto_secretstream_xchacha20poly1305_keygen();
        const { header, ciphertext } = encryptFileBody(plaintext, key);
        const thumbPush =
            sodium.crypto_secretstream_xchacha20poly1305_init_push(key);
        const file = buildMockEnteFile(key, header, thumbPush.header);
        file.metadata.title = "fallback-name.png";

        const api = new ApiClient({ fetch: mockFetchForBody(ciphertext) });
        const result = await downloadFile(api, file);

        expect(result.path).toBe("fallback-name.png");
        expect(readFileSync(result.path)).toEqual(Buffer.from(plaintext));
        // Clean up since it writes to cwd
        if (existsSync(result.path)) rmSync(result.path);
    });

    it("handles a larger single-chunk file (random binary payload)", async () => {
        // Most photos are under 4 MiB and therefore a single secretstream
        // chunk. This test exercises a non-trivial payload size with
        // random binary data (not just ASCII) to verify no encoding bugs.
        const plaintext = sodium.randombytes_buf(100_000);
        const key = sodium.crypto_secretstream_xchacha20poly1305_keygen();
        const { header, ciphertext } = encryptFileBody(plaintext, key);
        const thumbPush =
            sodium.crypto_secretstream_xchacha20poly1305_init_push(key);
        const file = buildMockEnteFile(key, header, thumbPush.header);

        const api = new ApiClient({ fetch: mockFetchForBody(ciphertext) });
        const outPath = join(testDir, "large-single.bin");
        const result = await downloadFile(api, file, outPath);

        expect(result.bytesWritten).toBe(100_000);
        expectSameBytes(readFileSync(outPath), plaintext);
    });

    it("decrypts a body that spans several secretstream chunks", async () => {
        // Files over 4 MiB arrive as several ciphertext chunks concatenated
        // into one HTTP body. The downloader has to re-split them on the
        // exact chunk boundary; getting that wrong corrupts every large
        // photo in an account. This is also the positive control for the
        // truncation tests below: it proves the multi-chunk fixture itself
        // decrypts cleanly when nothing has been removed from it.
        const { api, file } = fixtureFor(
            multiChunkKey,
            multiChunk.header,
            multiChunk.body,
        );
        const outPath = join(testDir, "multi-chunk.bin");

        const result = await downloadFile(api, file, outPath);

        expect(result.bytesWritten).toBe(multiChunk.plaintext.length);
        expectSameBytes(readFileSync(outPath), multiChunk.plaintext);
    });
});

describe("downloadThumbnail", () => {
    it("downloads and decrypts the thumbnail stream", async () => {
        const plaintext = new Uint8Array([0xff, 0xd8, 0xff, 0xe0]); // JPEG SOI
        const key = sodium.crypto_secretstream_xchacha20poly1305_keygen();

        const filePush =
            sodium.crypto_secretstream_xchacha20poly1305_init_push(key);
        const { header: thumbHeader, ciphertext: thumbCipher } =
            encryptFileBody(plaintext, key);

        const file = buildMockEnteFile(key, filePush.header, thumbHeader);
        const api = new ApiClient({ fetch: mockFetchForBody(thumbCipher) });
        const outPath = join(testDir, "thumb.jpg");

        const result = await downloadThumbnail(api, file, outPath);

        expect(result).toEqual({ path: outPath, bytesWritten: 4 });
        expect(readFileSync(outPath)).toEqual(Buffer.from(plaintext));
    });
});

// ---------------------------------------------------------------------------
// Truncation detection and atomic writes
//
// Everything below is the failure contract. It is deliberately written once
// per entry point via `entryPoints`, because `downloadFile` and
// `downloadThumbnail` must behave identically here: a corrupt thumbnail is
// just as unacceptable as a corrupt original, and `runBackup` trusts both.
// ---------------------------------------------------------------------------

describe.each(entryPoints)(
    "$name truncation handling",
    ({ name, download }) => {
        /** A fresh, empty directory so leftover-file assertions are meaningful. */
        const freshDir = (): string => {
            const dir = mkdtempSync(join(testDir, `${name}-`));
            return dir;
        };

        it("rejects a body whose final TAG_FINAL chunk never arrived", async () => {
            // Simulate a connection that dropped after the first 4 MiB chunk.
            // Every byte that did arrive decrypts and authenticates perfectly —
            // that is precisely the danger. The only signal that the file is
            // incomplete is the absence of a chunk tagged TAG_FINAL, so the
            // downloader must treat "stream ended on TAG_MESSAGE" as a hard
            // error rather than returning a short file.
            const truncatedBody = multiChunk.body.slice(
                0,
                multiChunk.finalChunkOffset,
            );
            const { api, file } = fixtureFor(
                multiChunkKey,
                multiChunk.header,
                truncatedBody,
            );
            const outPath = join(freshDir(), "truncated.bin");

            await expect(download(api, file, outPath)).rejects.toThrow(
                /truncated/i,
            );
        });

        it("rejects an empty body instead of writing a zero-byte file", async () => {
            // Ente always emits at least one chunk, even for empty content:
            // `encryptBlob` shows that a zero-length plaintext still produces a
            // TAG_FINAL chunk. A body with no chunks at all therefore means the
            // transfer failed, not that the file is empty. Writing a zero-byte
            // file here would be the worst outcome, because `runBackup` would
            // then see a file it considers present and never retry it.
            const { api, file } = fixtureFor(
                multiChunkKey,
                multiChunk.header,
                new Uint8Array(0),
            );
            const outPath = join(freshDir(), "empty.bin");

            await expect(download(api, file, outPath)).rejects.toThrow(
                /truncated/i,
            );
        });

        it("leaves no file at the destination after a truncated download", async () => {
            // The caller's contract: if the promise rejects, the destination
            // path does not exist. Nothing downstream should ever have to guess
            // whether a leftover file is complete.
            //
            // The body here is a single chunk that was never tagged TAG_FINAL,
            // which puts the downloader in exactly the terminal state a lost
            // last chunk produces, without the cost of a 4 MiB fixture. What
            // this test is really about is the state of the filesystem after
            // the rejection.
            const key = sodium.crypto_secretstream_xchacha20poly1305_keygen();
            const { header, ciphertext } = encryptNonFinalBody(
                sodium.randombytes_buf(256),
                key,
            );
            const { api, file } = fixtureFor(key, header, ciphertext);
            const dir = freshDir();
            const outPath = join(dir, "absent.bin");

            await expect(download(api, file, outPath)).rejects.toThrow(
                /truncated/i,
            );

            expect(existsSync(outPath)).toBe(false);
            // And no temporary scratch file is left behind either: the download
            // stages plaintext in a sibling temp file, which must be removed on
            // the failure path so repeated failures cannot fill the disk.
            expect(readdirSync(dir)).toEqual([]);
        });

        it("leaves no file at the destination when a chunk fails authentication", async () => {
            // The same guarantee has to hold for every failure mode, not just
            // truncation. Here a byte of ciphertext is flipped, so Poly1305
            // verification fails inside `pullStreamChunk`. The error must
            // propagate unchanged (it is the real diagnosis) and the destination
            // must still be untouched.
            const plaintext = sodium.randombytes_buf(256);
            const key = sodium.crypto_secretstream_xchacha20poly1305_keygen();
            const { header, ciphertext } = encryptFileBody(plaintext, key);
            const corrupted = Uint8Array.from(ciphertext);
            corrupted[10] ^= 0xff;

            const { api, file } = fixtureFor(key, header, corrupted);
            const dir = freshDir();
            const outPath = join(dir, "corrupt.bin");

            await expect(download(api, file, outPath)).rejects.toThrow(
                /authentication failed/i,
            );

            expect(existsSync(outPath)).toBe(false);
            expect(readdirSync(dir)).toEqual([]);
        });

        it("does not clobber an existing file when the download fails", async () => {
            // The repair case. A user re-running a backup over a directory that
            // already holds good originals must never end up worse off: a failed
            // download leaves the previous contents exactly as they were, so the
            // old good copy survives until a complete new one is available to
            // replace it in a single rename.
            const existing = new TextEncoder().encode(
                "previously downloaded, known-good contents",
            );
            const key = sodium.crypto_secretstream_xchacha20poly1305_keygen();
            const { header, ciphertext } = encryptNonFinalBody(
                sodium.randombytes_buf(256),
                key,
            );
            const { api, file } = fixtureFor(key, header, ciphertext);
            const dir = freshDir();
            const outPath = join(dir, "existing.bin");
            writeFileSync(outPath, existing);

            await expect(download(api, file, outPath)).rejects.toThrow(
                /truncated/i,
            );

            expect(readFileSync(outPath)).toEqual(Buffer.from(existing));
            expect(readdirSync(dir)).toEqual(["existing.bin"]);
        });

        it("replaces an existing file when the download succeeds", async () => {
            // The mirror image of the previous test: a complete download does
            // overwrite whatever was at the destination, atomically, via rename.
            const existing = new TextEncoder().encode("stale contents");
            const plaintext = sodium.randombytes_buf(512);
            const key = sodium.crypto_secretstream_xchacha20poly1305_keygen();
            const { header, ciphertext } = encryptFileBody(plaintext, key);

            const { api, file } = fixtureFor(key, header, ciphertext);
            const dir = freshDir();
            const outPath = join(dir, "replaced.bin");
            writeFileSync(outPath, existing);

            const result = await download(api, file, outPath);

            expect(result).toEqual({ path: outPath, bytesWritten: 512 });
            expect(readFileSync(outPath)).toEqual(Buffer.from(plaintext));
            // The temp file is gone once the rename has happened, so a
            // successful download leaves exactly one file behind.
            expect(readdirSync(dir)).toEqual(["replaced.bin"]);
        });
    },
);