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`. The last
 * chunk carries `TAG_FINAL`; any truncation is detected because the tag
 * will be missing.
 *
 * 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, readFileSync, rmSync, mkdirSync } from "node:fs";
import { join } from "node:path";
import { tmpdir } from "node:os";
import sodium from "libsodium-wrappers-sumo";
import { beforeAll, afterAll, describe, expect, it } from "vitest";
import { init, toBase64 } 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 = join(tmpdir(), `quack-test-${Date.now()}`);
    mkdirSync(testDir, { recursive: true });
});

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 `plaintext` as a multi-chunk secretstream. Each chunk is at
 * most `chunkSize` bytes of plaintext.
 */
const encryptMultiChunk = (
    plaintext: Uint8Array,
    key: Uint8Array,
    chunkSize: number,
): { header: Uint8Array; ciphertext: Uint8Array } => {
    const push = sodium.crypto_secretstream_xchacha20poly1305_init_push(key);
    const chunks: Uint8Array[] = [];
    for (let offset = 0; offset < plaintext.length; offset += chunkSize) {
        const end = Math.min(offset + chunkSize, plaintext.length);
        const isLast = end === plaintext.length;
        const tag = isLast
            ? sodium.crypto_secretstream_xchacha20poly1305_TAG_FINAL
            : sodium.crypto_secretstream_xchacha20poly1305_TAG_MESSAGE;
        chunks.push(
            sodium.crypto_secretstream_xchacha20poly1305_push(
                push.state,
                plaintext.slice(offset, end),
                null,
                tag,
            ),
        );
    }
    const total = chunks.reduce((n, c) => n + c.length, 0);
    const ciphertext = new Uint8Array(total);
    let pos = 0;
    for (const c of chunks) {
        ciphertext.set(c, pos);
        pos += c.length;
    }
    return { header: push.header, ciphertext };
};

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

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

describe("downloadFile", () => {
    it("downloads, decrypts, and writes a single-chunk file", async () => {
        const plaintext = new TextEncoder().encode(
            "Hello from quack! 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);

        expect(result.path).toBe(outPath);
        expect(result.bytesWritten).toBe(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");
        // Clean up since it writes to cwd
        if (existsSync(result.path)) rmSync(result.path);
    });

    it("handles a multi-chunk encrypted stream", async () => {
        // Encrypt with a small chunk size so we get multiple chunks
        // without allocating megabytes. The download function uses the
        // real STREAM_CHUNK_SIZE (4 MiB) boundary to split, but when
        // the total ciphertext is smaller than one full chunk, the
        // entire body is the "last" chunk. To test multi-chunk, we
        // encrypt with 4 MiB chunks and a body just over that boundary.
        //
        // For unit test speed, we use a 50-byte "chunk size" to produce
        // multiple chunks, then concatenate the ciphertext. The download
        // function must still decrypt correctly because the chunk
        // boundaries are determined by the secretstream tag bytes, and
        // Ente's convention is that all chunks except the last are
        // exactly STREAM_CHUNK_SIZE + STREAM_CHUNK_OVERHEAD bytes.
        //
        // In practice, files under 4 MiB are single-chunk (the common
        // case). The multi-chunk path is exercised by the live
        // integration test against large photos.
        const plaintext = sodium.randombytes_buf(200);
        const key = sodium.crypto_secretstream_xchacha20poly1305_keygen();
        // Use real 4 MiB chunk size for encryption to match what Ente
        // produces. Since plaintext is 200 bytes, this will be a single
        // chunk anyway, but let's also test with the multi-chunk helper.
        const { header, ciphertext } = encryptMultiChunk(plaintext, key, 50);
        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, "multi-chunk.bin");
        const result = await downloadFile(api, file, outPath);

        expect(result.bytesWritten).toBe(200);
        expect(readFileSync(outPath)).toEqual(Buffer.from(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.bytesWritten).toBe(4);
        expect(readFileSync(outPath)).toEqual(Buffer.from(plaintext));
    });
});