quak/test/thumbnails/thumbnails.test.ts

/**
 * Tests for `listMissingThumbnails` and `fixMissingThumbnails`.
 *
 * These use a mock Ente server that serves encrypted collections, files,
 * and thumbnails. The mock server has deliberate gaps: some files have
 * working thumbnails, others return 404 or empty bodies. The tests
 * verify that the detection and repair logic handles each case correctly.
 *
 * `fixMissingThumbnails` is the most complex function in quak: it
 * downloads the original file, generates a JPEG thumbnail with jpeg-js,
 * encrypts it with secretstream push, gets a presigned upload URL,
 * uploads to S3, and registers the new thumbnail with the API. The
 * test verifies each step actually happened and the uploaded data is
 * a valid encrypted blob that decrypts to a JPEG.
 */

import sodium from "libsodium-wrappers-sumo";
import * as jpegJs from "jpeg-js";
import { beforeAll, describe, expect, it } from "vitest";
import {
    init,
    toBase64,
    decryptBlob,
    fromBase64,
    deriveKEK,
    deriveLoginSubkey,
} from "../../src/crypto/index.js";
import { SRP, SrpServer } from "fast-srp-hap";
import { Client } from "../../src/client.js";
import {
    listMissingThumbnails,
    fixMissingThumbnails,
} from "../../src/thumbnails.js";
import type { KeyAttributes } from "../../src/auth/types.js";

// ---------------------------------------------------------------------------
// Mock server with controllable thumbnail behavior
// ---------------------------------------------------------------------------

const TEST_EMAIL = "thumb@example.com";
const TEST_PASSWORD = "thumbpass";
const TEST_OPS = 2;
const TEST_MEM = 64 * 1024 * 1024;

interface ThumbMockState {
    verifier: Buffer;
    srpAttributes: Record<string, unknown>;
    keyAttributes: KeyAttributes;
    encryptedToken: string;
    collections: Record<string, unknown>[];
    filesByCollection: Record<number, Record<string, unknown>[]>;
    fileCiphertexts: Record<number, Uint8Array>;
    fileKeys: Record<number, Uint8Array>;
    thumbnailBehavior: Record<number, "ok" | "empty" | "404">;
    // Captures from fix operations
    uploadedThumbnails: {
        fileID: number;
        objectKey: string;
        decryptionHeader: string;
        ciphertext: Uint8Array;
    }[];
}

let mock: ThumbMockState;

const buildThumbMock = async (): Promise<ThumbMockState> => {
    const kekSalt = sodium.randombytes_buf(sodium.crypto_pwhash_SALTBYTES);
    const kek = await deriveKEK(TEST_PASSWORD, kekSalt, TEST_OPS, TEST_MEM);
    const loginSubKeyBytes = deriveLoginSubkey(kek);

    const srpUserID = "thumb-srp";
    const srpSalt = sodium.randombytes_buf(16);
    const verifier = SRP.computeVerifier(
        SRP.params["4096"],
        Buffer.from(srpSalt),
        Buffer.from(srpUserID),
        Buffer.from(loginSubKeyBytes),
    );

    const masterKey = sodium.randombytes_buf(32);
    const keyNonce = sodium.randombytes_buf(sodium.crypto_secretbox_NONCEBYTES);
    const encryptedKey = sodium.crypto_secretbox_easy(masterKey, keyNonce, kek);
    const kp = sodium.crypto_box_keypair();
    const skNonce = sodium.randombytes_buf(sodium.crypto_secretbox_NONCEBYTES);
    const encSK = sodium.crypto_secretbox_easy(
        kp.privateKey,
        skNonce,
        masterKey,
    );
    const tokenBytes = sodium.randombytes_buf(32);
    const encToken = sodium.crypto_box_seal(tokenBytes, kp.publicKey);

    const keyAttributes: KeyAttributes = {
        kekSalt: toBase64(kekSalt),
        encryptedKey: toBase64(encryptedKey),
        keyDecryptionNonce: toBase64(keyNonce),
        publicKey: toBase64(kp.publicKey),
        encryptedSecretKey: toBase64(encSK),
        secretKeyDecryptionNonce: toBase64(skNonce),
        memLimit: TEST_MEM,
        opsLimit: TEST_OPS,
    };

    // One collection with 3 files: ok thumbnail, empty thumbnail, 404 thumbnail
    const collKey = sodium.crypto_secretbox_keygen();
    const ckN = sodium.randombytes_buf(sodium.crypto_secretbox_NONCEBYTES);
    const encCK = sodium.crypto_secretbox_easy(collKey, ckN, masterKey);
    const nameB = new TextEncoder().encode("Photos");
    const cnN = sodium.randombytes_buf(sodium.crypto_secretbox_NONCEBYTES);
    const encCN = sodium.crypto_secretbox_easy(nameB, cnN, collKey);

    const rawCollection = {
        id: 1,
        owner: { id: 42 },
        encryptedKey: toBase64(encCK),
        keyDecryptionNonce: toBase64(ckN),
        encryptedName: toBase64(encCN),
        nameDecryptionNonce: toBase64(cnN),
        type: "album",
        updationTime: 1700000000000000,
    };

    // Generate a real tiny JPEG via jpeg-js
    const w = 100;
    const h = 80;
    const pixels = new Uint8Array(w * h * 4);
    for (let i = 0; i < pixels.length; i += 4) {
        pixels[i] = 255; // R
        pixels[i + 1] = 0; // G
        pixels[i + 2] = 0; // B
        pixels[i + 3] = 255; // A
    }
    const tinyJpeg = jpegJs.encode(
        { data: pixels, width: w, height: h },
        80,
    ).data;

    const fileKeys: Record<number, Uint8Array> = {};
    const fileCiphertexts: Record<number, Uint8Array> = {};
    const rawFiles: Record<string, unknown>[] = [];

    for (const fileID of [100, 101, 102]) {
        const fk = sodium.crypto_secretstream_xchacha20poly1305_keygen();
        fileKeys[fileID] = fk;
        const fkN = sodium.randombytes_buf(sodium.crypto_secretbox_NONCEBYTES);
        const encFK = sodium.crypto_secretbox_easy(fk, fkN, collKey);

        const meta = JSON.stringify({
            title: `file-${fileID}.jpg`,
            fileType: 0,
            creationTime: 1700000000000000,
            modificationTime: 1700000000000000,
        });
        const metaPush =
            sodium.crypto_secretstream_xchacha20poly1305_init_push(fk);
        const encMeta = sodium.crypto_secretstream_xchacha20poly1305_push(
            metaPush.state,
            new TextEncoder().encode(meta),
            null,
            sodium.crypto_secretstream_xchacha20poly1305_TAG_FINAL,
        );

        // Encrypt the tiny JPEG as the file body
        const filePush =
            sodium.crypto_secretstream_xchacha20poly1305_init_push(fk);
        const encFile = sodium.crypto_secretstream_xchacha20poly1305_push(
            filePush.state,
            new Uint8Array(tinyJpeg),
            null,
            sodium.crypto_secretstream_xchacha20poly1305_TAG_FINAL,
        );
        fileCiphertexts[fileID] = encFile;

        rawFiles.push({
            id: fileID,
            collectionID: 1,
            ownerID: 42,
            encryptedKey: toBase64(encFK),
            keyDecryptionNonce: toBase64(fkN),
            metadata: {
                encryptedData: toBase64(encMeta),
                decryptionHeader: toBase64(metaPush.header),
            },
            file: { decryptionHeader: toBase64(filePush.header) },
            thumbnail: {
                decryptionHeader: toBase64(sodium.randombytes_buf(24)),
            },
            updationTime: 1700000000000000,
        });
    }

    return {
        verifier,
        srpAttributes: {
            srpUserID,
            srpSalt: toBase64(srpSalt),
            memLimit: TEST_MEM,
            opsLimit: TEST_OPS,
            kekSalt: toBase64(kekSalt),
            isEmailMFAEnabled: false,
        },
        keyAttributes,
        encryptedToken: toBase64(encToken),
        collections: [rawCollection],
        filesByCollection: { 1: rawFiles },
        fileCiphertexts,
        fileKeys,
        thumbnailBehavior: {
            100: "ok",
            101: "empty",
            102: "404",
        },
        uploadedThumbnails: [],
    };
};

const buildThumbFetch = (m: ThumbMockState) => {
    let srpServer: SrpServer;

    return (async (
        input: RequestInfo | URL,
        init?: RequestInit,
    ): Promise<Response> => {
        const url =
            typeof input === "string"
                ? input
                : input instanceof URL
                  ? input.href
                  : input.url;
        const parsed = new URL(url);
        const path = parsed.pathname;
        const json = (body: unknown) =>
            new Response(JSON.stringify(body), {
                status: 200,
                headers: { "content-type": "application/json" },
            });

        // SRP auth flow
        if (path === "/users/srp/attributes")
            return json({ attributes: m.srpAttributes });
        if (path === "/users/srp/create-session") {
            const body = JSON.parse(init?.body as string);
            const serverKey = await SRP.genKey();
            srpServer = new SrpServer(
                SRP.params["4096"],
                m.verifier,
                serverKey,
            );
            const B = srpServer.computeB();
            srpServer.setA(Buffer.from(body.srpA, "base64"));
            return json({ sessionID: "s1", srpB: B.toString("base64") });
        }
        if (path === "/users/srp/verify-session") {
            const body = JSON.parse(init?.body as string);
            srpServer.checkM1(Buffer.from(body.srpM1, "base64"));
            return json({
                srpM2: srpServer.computeM2().toString("base64"),
                id: 42,
                keyAttributes: m.keyAttributes,
                encryptedToken: m.encryptedToken,
            });
        }

        // Collections & files
        if (path === "/collections/v2")
            return json({ collections: m.collections });
        if (path === "/collections/v2/diff") {
            const collID = Number(parsed.searchParams.get("collectionID"));
            return json({
                diff: m.filesByCollection[collID] ?? [],
                hasMore: false,
            });
        }

        // File download (for fix: download original to generate thumb)
        if (
            url.includes("files.ente.io") ||
            path.startsWith("/files/download/")
        ) {
            const fileID = Number(
                parsed.searchParams.get("fileID") ?? path.split("/").pop(),
            );
            const ct = m.fileCiphertexts[fileID];
            if (ct) return new Response(ct, { status: 200 });
            return new Response("not found", { status: 404 });
        }

        // Thumbnail download (for list: check if thumbnail exists)
        if (
            url.includes("thumbnails.ente.io") ||
            path.startsWith("/files/preview/")
        ) {
            const fileID = Number(
                parsed.searchParams.get("fileID") ?? path.split("/").pop(),
            );
            const behavior = m.thumbnailBehavior[fileID];
            if (behavior === "ok") {
                return new Response(new Uint8Array([0xff, 0xd8, 0xff]), {
                    status: 200,
                });
            }
            if (behavior === "empty") {
                return new Response(new Uint8Array(0), { status: 200 });
            }
            return new Response("not found", { status: 404 });
        }

        // Upload URL minting
        if (path === "/files/upload-url") {
            return json({
                objectKey: `42/thumb-${Date.now()}`,
                url: "https://s3.mock.test/presigned-put",
            });
        }

        // Presigned PUT (S3 upload)
        if (url.startsWith("https://s3.mock.test/")) {
            const body = init?.body;
            // Store the uploaded bytes for later inspection
            if (body instanceof Uint8Array) {
                (m as Record<string, unknown>)._lastUploadedCiphertext = body;
            } else if (body instanceof ArrayBuffer) {
                (m as Record<string, unknown>)._lastUploadedCiphertext =
                    new Uint8Array(body);
            }
            return new Response(null, { status: 200 });
        }

        // Update thumbnail metadata
        if (path === "/files/thumbnail" && init?.method === "PUT") {
            const reqBody = JSON.parse(init?.body as string);
            m.uploadedThumbnails.push({
                fileID: reqBody.fileID,
                objectKey: reqBody.thumbnail.objectKey,
                decryptionHeader: reqBody.thumbnail.decryptionHeader,
                ciphertext:
                    ((m as Record<string, unknown>)
                        ._lastUploadedCiphertext as Uint8Array) ??
                    new Uint8Array(0),
            });
            return json({});
        }

        return new Response("not found", { status: 404 });
    }) as typeof globalThis.fetch;
};

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

beforeAll(async () => {
    await init();
    await sodium.ready;
    mock = await buildThumbMock();
});

describe("listMissingThumbnails", () => {
    it("identifies files with empty and 404 thumbnails, ignores working ones", async () => {
        const client = await Client.login({
            email: TEST_EMAIL,
            password: TEST_PASSWORD,
            apiOptions: { fetch: buildThumbFetch(mock) },
        });

        const missing = await listMissingThumbnails(client);

        // File 100 has a working thumbnail → not reported
        // File 101 has an empty thumbnail → reported
        // File 102 has a 404 thumbnail → reported
        expect(missing.length).toBe(2);
        const ids = missing.map((m) => m.fileID).sort();
        expect(ids).toEqual([101, 102]);

        const emptyEntry = missing.find((m) => m.fileID === 101)!;
        expect(emptyEntry.reason).toContain("empty");
        expect(emptyEntry.title).toBe("file-101.jpg");
        expect(emptyEntry.collection).toBe("Photos");

        const notFoundEntry = missing.find((m) => m.fileID === 102)!;
        expect(notFoundEntry.reason).toContain("fetch failed");
    });

    it("deduplicates files seen in multiple collections", async () => {
        // Add the same files to a second collection in the mock
        const mockWithDupes = await buildThumbMock();
        const dupeCollection = {
            ...mockWithDupes.collections[0]!,
            id: 2,
        };
        mockWithDupes.collections.push(
            dupeCollection as Record<string, unknown>,
        );
        mockWithDupes.filesByCollection[2] =
            mockWithDupes.filesByCollection[1]!;

        const client = await Client.login({
            email: TEST_EMAIL,
            password: TEST_PASSWORD,
            apiOptions: { fetch: buildThumbFetch(mockWithDupes) },
        });

        const missing = await listMissingThumbnails(client);

        // Should still be 2, not 4 (each file checked only once)
        expect(missing.length).toBe(2);
    });
});

describe("fixMissingThumbnails", () => {
    it("downloads original, generates thumbnail, encrypts, uploads, and registers", async () => {
        const fixMock = await buildThumbMock();
        const client = await Client.login({
            email: TEST_EMAIL,
            password: TEST_PASSWORD,
            apiOptions: { fetch: buildThumbFetch(fixMock) },
        });

        const results = await fixMissingThumbnails(client, [101]);

        expect(results.length).toBe(1);
        expect(results[0]!.success).toBe(true);
        expect(results[0]!.fileID).toBe(101);
        expect(results[0]!.title).toBe("file-101.jpg");
        expect(results[0]!.collection).toBe("Photos");

        // Verify the uploaded thumbnail was registered
        expect(fixMock.uploadedThumbnails.length).toBe(1);
        const upload = fixMock.uploadedThumbnails[0]!;
        expect(upload.fileID).toBe(101);
        expect(upload.objectKey).toContain("42/");
        expect(upload.decryptionHeader.length).toBeGreaterThan(0);

        // Verify the uploaded ciphertext can be decrypted back to a JPEG
        const fileKey = fixMock.fileKeys[101]!;
        const decrypted = decryptBlob(
            upload.ciphertext,
            fromBase64(upload.decryptionHeader),
            fileKey,
        );
        // JPEG magic bytes: FF D8 FF
        expect(decrypted[0]).toBe(0xff);
        expect(decrypted[1]).toBe(0xd8);
        expect(decrypted[2]).toBe(0xff);
        // Verify jpeg-js produced a reasonably sized thumbnail
        expect(decrypted.length).toBeGreaterThan(100);
        expect(decrypted.length).toBeLessThan(50000);
    });

    it("reports failure for nonexistent file IDs without crashing", async () => {
        const fixMock = await buildThumbMock();
        const client = await Client.login({
            email: TEST_EMAIL,
            password: TEST_PASSWORD,
            apiOptions: { fetch: buildThumbFetch(fixMock) },
        });

        const results = await fixMissingThumbnails(client, [999]);

        expect(results.length).toBe(1);
        expect(results[0]!.success).toBe(false);
        expect(results[0]!.fileID).toBe(999);
        expect(results[0]!.error).toContain("not found");
    });

    it("continues after one file fails and reports mixed results", async () => {
        const fixMock = await buildThumbMock();
        // Make file 102 fail by removing its ciphertext so download fails
        delete fixMock.fileCiphertexts[102];

        const client = await Client.login({
            email: TEST_EMAIL,
            password: TEST_PASSWORD,
            apiOptions: { fetch: buildThumbFetch(fixMock) },
        });

        const results = await fixMissingThumbnails(client, [101, 102]);

        expect(results.length).toBe(2);
        const success = results.find((r) => r.fileID === 101)!;
        const failure = results.find((r) => r.fileID === 102)!;
        expect(success.success).toBe(true);
        expect(failure.success).toBe(false);
    });
});

describe("Client.getApiClient", () => {
    it("returns the ApiClient when logged in", async () => {
        const client = await Client.login({
            email: TEST_EMAIL,
            password: TEST_PASSWORD,
            apiOptions: { fetch: buildThumbFetch(mock) },
        });

        const api = client.getApiClient();
        expect(api).toBeDefined();
        expect(typeof api.getJSON).toBe("function");
    });

    it("throws after logout", async () => {
        const client = await Client.login({
            email: TEST_EMAIL,
            password: TEST_PASSWORD,
            apiOptions: { fetch: buildThumbFetch(mock) },
        });
        client.logout();

        expect(() => client.getApiClient()).toThrow(/logged out/);
    });
});