SnowLuma/packages/onebot/tests/instance-rkey.test.ts

import { describe, it, expect, vi, afterEach } from 'vitest';
import { RKeyCache } from '../src/instance-rkey';
import type { MessageElement } from '@snowluma/protocol/events';

// The screenshot URL: a raw QQ-NT multimedia download link with no rkey yet.
// appid 1407 = group image, 1406 = private/c2c image (see instance-rkey).
const NT_IMAGE_URL = 'https://multimedia.nt.qq.com.cn/download?appid=1407&fileid=ABC';
const NT_PRIVATE_IMAGE_URL = 'https://multimedia.nt.qq.com.cn/download?appid=1406&fileid=XYZ';
const GROUP_IMAGE_TYPE = 20;
const PRIVATE_IMAGE_TYPE = 10;

function nowSec(): number {
    return Math.floor(Date.now() / 1000);
}

function makeBridge(fetchImpl: () => Promise<unknown[]>) {
    const fetchDownloadRKeys = vi.fn(fetchImpl);
    return { bridge: { apis: { contacts: { fetchDownloadRKeys } } }, fetchDownloadRKeys };
}

function freshGroupRkey(value = 'GROUPKEY') {
    return [
        { rkey: value, type: GROUP_IMAGE_TYPE, ttlSeconds: 3600, createTime: nowSec(), storeId: 0 },
    ];
}

function freshBothRkeys() {
    return [
        {
            rkey: 'PRIVATEKEY',
            type: PRIVATE_IMAGE_TYPE,
            ttlSeconds: 3600,
            createTime: nowSec(),
            storeId: 0,
        },
        {
            rkey: 'GROUPKEY',
            type: GROUP_IMAGE_TYPE,
            ttlSeconds: 3600,
            createTime: nowSec(),
            storeId: 0,
        },
    ];
}

function imageEl(): MessageElement {
    return { type: 'image', imageUrl: NT_IMAGE_URL };
}

function privateImageEl(): MessageElement {
    return { type: 'image', imageUrl: NT_PRIVATE_IMAGE_URL };
}

describe('RKeyCache.resolveImageUrl', () => {
    it('awaits a refresh on a cold cache and appends the rkey', async () => {
        const { bridge, fetchDownloadRKeys } = makeBridge(async () => freshGroupRkey());
        const cache = new RKeyCache();

        const url = await cache.resolveImageUrl(bridge as never, imageEl(), true);

        expect(url).toBe(`${NT_IMAGE_URL}&rkey=GROUPKEY`);
        expect(fetchDownloadRKeys).toHaveBeenCalledTimes(1);
    });

    it('coalesces concurrent cold-cache lookups into a single fetch', async () => {
        const { bridge, fetchDownloadRKeys } = makeBridge(async () => freshGroupRkey());
        const cache = new RKeyCache();

        const urls = await Promise.all([
            cache.resolveImageUrl(bridge as never, imageEl(), true),
            cache.resolveImageUrl(bridge as never, imageEl(), true),
            cache.resolveImageUrl(bridge as never, imageEl(), true),
        ]);

        expect(urls).toEqual(Array(3).fill(`${NT_IMAGE_URL}&rkey=GROUPKEY`));
        expect(fetchDownloadRKeys).toHaveBeenCalledTimes(1);
    });

    it('serves a still-valid rkey from cache without re-fetching', async () => {
        const { bridge, fetchDownloadRKeys } = makeBridge(async () => freshGroupRkey());
        const cache = new RKeyCache();

        await cache.resolveImageUrl(bridge as never, imageEl(), true);
        await cache.resolveImageUrl(bridge as never, imageEl(), true);

        expect(fetchDownloadRKeys).toHaveBeenCalledTimes(1);
    });

    it('falls back to the unsigned URL (no throw) when the refresh fails', async () => {
        const { bridge, fetchDownloadRKeys } = makeBridge(async () => {
            throw new Error('oidb down');
        });
        const cache = new RKeyCache();

        const url = await cache.resolveImageUrl(bridge as never, imageEl(), true);

        expect(url).toBe(NT_IMAGE_URL);
        expect(fetchDownloadRKeys).toHaveBeenCalledTimes(1);
    });

    it('uses the private rkey type for c2c images (appid 1406)', async () => {
        const { bridge } = makeBridge(async () => freshBothRkeys());
        const cache = new RKeyCache();

        const url = await cache.resolveImageUrl(bridge as never, privateImageEl(), false);

        expect(url).toBe(`${NT_PRIVATE_IMAGE_URL}&rkey=PRIVATEKEY`);
    });

    // The scene follows the image's own appid, not the carrying message — the
    // crux of `/get_forward_msg` rkey correctness (issue #74): a c2c image
    // forwarded into a group must still be signed with the private rkey.
    it('signs an appid-1406 image with the private rkey even when isGroup is true', async () => {
        const { bridge } = makeBridge(async () => freshBothRkeys());
        const cache = new RKeyCache();

        const url = await cache.resolveImageUrl(bridge as never, privateImageEl(), true);

        expect(url).toBe(`${NT_PRIVATE_IMAGE_URL}&rkey=PRIVATEKEY`);
    });

    it('signs an appid-1407 image with the group rkey even when isGroup is false', async () => {
        const { bridge } = makeBridge(async () => freshBothRkeys());
        const cache = new RKeyCache();

        const url = await cache.resolveImageUrl(bridge as never, imageEl(), false);

        expect(url).toBe(`${NT_IMAGE_URL}&rkey=GROUPKEY`);
    });

    it('falls back to the isGroup scene when the URL carries no appid', async () => {
        const noAppidUrl = 'https://multimedia.nt.qq.com.cn/download?fileid=NOAPPID';
        const { bridge } = makeBridge(async () => freshBothRkeys());
        const cache = new RKeyCache();

        const url = await cache.resolveImageUrl(
            bridge as never,
            { type: 'image', imageUrl: noAppidUrl },
            true,
        );

        expect(url).toBe(`${noAppidUrl}&rkey=GROUPKEY`);
    });

    // #156: a transient empty/failed fetch must self-heal in seconds, not 30s.
    // An empty cache uses the short RKEY_EMPTY_COOLDOWN (10s), so a retry ~11s
    // later refetches — under the 30s steady-state backoff it would not.
    it('retries before the 30s cooldown while the cache holds no usable rkey', async () => {
        vi.useFakeTimers({ toFake: ['Date'] });

        try {
            vi.setSystemTime(new Date(1_700_000_000_000));
            let calls = 0;
            const { bridge, fetchDownloadRKeys } = makeBridge(async () =>
                ++calls === 1 ? [] : freshGroupRkey(),
            );
            const cache = new RKeyCache();

            const url1 = await cache.resolveImageUrl(bridge as never, imageEl(), true);
            expect(url1).toBe(NT_IMAGE_URL); // empty fetch → bare URL

            vi.setSystemTime(new Date(1_700_000_000_000 + 11_000)); // +11s: past 10s empty-cooldown, inside 30s
            const url2 = await cache.resolveImageUrl(bridge as never, imageEl(), true);

            expect(url2).toBe(`${NT_IMAGE_URL}&rkey=GROUPKEY`);
            expect(fetchDownloadRKeys).toHaveBeenCalledTimes(2);
        } finally {
            vi.useRealTimers();
        }
    });

    it('does NOT refetch within 30s once a usable rkey is cached', async () => {
        vi.useFakeTimers({ toFake: ['Date'] });

        try {
            vi.setSystemTime(new Date(1_700_000_000_000));
            const { bridge, fetchDownloadRKeys } = makeBridge(async () => freshGroupRkey());
            const cache = new RKeyCache();

            await cache.resolveImageUrl(bridge as never, imageEl(), true);
            vi.setSystemTime(new Date(1_700_000_000_000 + 15_000)); // +15s, inside 30s backoff
            await cache.resolveImageUrl(bridge as never, imageEl(), true);

            expect(fetchDownloadRKeys).toHaveBeenCalledTimes(1);
        } finally {
            vi.useRealTimers();
        }
    });

    it('leaves non-NT URLs untouched and never fetches', async () => {
        const { bridge, fetchDownloadRKeys } = makeBridge(async () => freshGroupRkey());
        const cache = new RKeyCache();
        const el: MessageElement = { type: 'image', imageUrl: 'https://gchat.qpic.cn/foo.jpg' };

        const url = await cache.resolveImageUrl(bridge as never, el, true);

        expect(url).toBe('https://gchat.qpic.cn/foo.jpg');
        expect(fetchDownloadRKeys).not.toHaveBeenCalled();
    });
});

describe('RKeyCache.resolveMediaUrl', () => {
    it.each([
        ['video', false, 12],
        ['video', true, 22],
        ['record', false, 14],
        ['record', true, 24],
        ['file', false, 10],
        ['file', true, 20],
    ] as const)('signs %s media (group=%s) with rkey type %s', async (type, isGroup, expected) => {
        const { bridge } = makeBridge(async () =>
            [10, 20, 12, 22, 14, 24].map((type) => ({
                rkey: `?rkey=key/${type}`,
                type,
                ttlSeconds: 3600,
                createTime: nowSec(),
            })),
        );
        const cache = new RKeyCache();
        const url = 'https://multimedia.nt.qq.com.cn/download?fileid=MEDIA';
        expect(await cache.resolveMediaUrl(bridge as never, { type, url }, isGroup)).toBe(
            `${url}&rkey=key%2F${expected}`,
        );
    });

    it('preserves signed, unrelated, and absent URLs without fetching', async () => {
        const { bridge, fetchDownloadRKeys } = makeBridge(async () => freshGroupRkey());
        const cache = new RKeyCache();

        for (const url of ['', 'https://example.com/video', `${NT_IMAGE_URL}&rkey=existing`]) {
            expect(await cache.resolveMediaUrl(bridge as never, { type: 'video', url }, true)).toBe(
                url,
            );
        }

        expect(fetchDownloadRKeys).not.toHaveBeenCalled();
    });
});

describe('RKeyCache remote fallback (#156)', () => {
    const SERVER = 'https://rkey.example/r';

    function stubFetch(payload: unknown, status = 200) {
        return vi
            .spyOn(globalThis, 'fetch')
            .mockResolvedValue(new Response(JSON.stringify(payload), { status }));
    }

    afterEach(() => {
        vi.restoreAllMocks();
    });

    it('is off by default: never contacts a server, leaves the URL bare', async () => {
        const fetchSpy = vi.spyOn(globalThis, 'fetch');
        const { bridge } = makeBridge(async () => []); // native yields nothing
        const cache = new RKeyCache(); // no config → fallback off

        const url = await cache.resolveImageUrl(bridge as never, imageEl(), true);

        expect(url).toBe(NT_IMAGE_URL);
        expect(fetchSpy).not.toHaveBeenCalled();
    });

    it('falls back to the configured server when the native fetch is empty', async () => {
        const fetchSpy = stubFetch({
            group_rkey: '&rkey=FBGROUP',
            private_rkey: '&rkey=FBPRIV',
            expired_time: nowSec() + 3600,
        });
        const { bridge } = makeBridge(async () => []);
        const cache = new RKeyCache({ fallbackServers: [SERVER] });

        const url = await cache.resolveImageUrl(bridge as never, imageEl(), true);

        expect(url).toBe(`${NT_IMAGE_URL}&rkey=FBGROUP`);
        expect(fetchSpy).toHaveBeenCalledWith(SERVER, expect.anything());
    });

    it('signs c2c images with the fallback private rkey', async () => {
        stubFetch({
            group_rkey: '&rkey=FBGROUP',
            private_rkey: '&rkey=FBPRIV',
            expired_time: nowSec() + 3600,
        });
        const { bridge } = makeBridge(async () => []);
        const cache = new RKeyCache({ fallbackServers: [SERVER] });

        const url = await cache.resolveImageUrl(bridge as never, privateImageEl(), false);

        expect(url).toBe(`${NT_PRIVATE_IMAGE_URL}&rkey=FBPRIV`);
    });

    it('accepts an OneBot { retcode, data } wrapper and bare (unprefixed) tokens', async () => {
        stubFetch({
            retcode: 0,
            data: { group_rkey: 'BAREKEY', private_rkey: '', expired_time: nowSec() + 3600 },
        });
        const { bridge } = makeBridge(async () => []);
        const cache = new RKeyCache({ fallbackServers: [SERVER] });

        const url = await cache.resolveImageUrl(bridge as never, imageEl(), true);

        expect(url).toBe(`${NT_IMAGE_URL}&rkey=BAREKEY`);
    });

    it('does not hit any server when the native fetch already returned a rkey', async () => {
        const fetchSpy = stubFetch({ group_rkey: '&rkey=FBGROUP', expired_time: nowSec() + 3600 });
        const { bridge } = makeBridge(async () => freshGroupRkey());
        const cache = new RKeyCache({ fallbackServers: [SERVER] });

        const url = await cache.resolveImageUrl(bridge as never, imageEl(), true);

        expect(url).toBe(`${NT_IMAGE_URL}&rkey=GROUPKEY`); // native rkey, not fallback
        expect(fetchSpy).not.toHaveBeenCalled();
    });

    it('tries the next server when the first one fails', async () => {
        const fetchSpy = vi
            .spyOn(globalThis, 'fetch')
            .mockRejectedValueOnce(new Error('down'))
            .mockResolvedValueOnce(
                new Response(
                    JSON.stringify({ group_rkey: '&rkey=SECOND', expired_time: nowSec() + 3600 }),
                    { status: 200 },
                ),
            );
        const { bridge } = makeBridge(async () => []);
        const cache = new RKeyCache({
            fallbackServers: ['https://a.example/r', 'https://b.example/r'],
        });

        const url = await cache.resolveImageUrl(bridge as never, imageEl(), true);

        expect(url).toBe(`${NT_IMAGE_URL}&rkey=SECOND`);
        expect(fetchSpy).toHaveBeenCalledTimes(2);
    });

    it('leaves the URL bare when every fallback server fails', async () => {
        vi.spyOn(globalThis, 'fetch').mockRejectedValue(new Error('down'));
        const { bridge } = makeBridge(async () => []);
        const cache = new RKeyCache({ fallbackServers: [SERVER] });

        const url = await cache.resolveImageUrl(bridge as never, imageEl(), true);

        expect(url).toBe(NT_IMAGE_URL);
    });

    // Hot-reload wire: manager.reloadGlobalSettings → instance.applyGlobalSettings
    // → rkeyCache.setFallbackServers. Verify the new list actually takes effect.
    // (Fake Date so the 2nd lookup clears the empty-cache refresh cooldown.)
    it('setFallbackServers turns the fallback on for a cache that started with none', async () => {
        vi.useFakeTimers({ toFake: ['Date'] });

        try {
            vi.setSystemTime(new Date(1_700_000_000_000));
            const fetchSpy = stubFetch({
                group_rkey: '&rkey=HOTGROUP',
                expired_time: 1_700_000_000 + 3600,
            });
            const { bridge } = makeBridge(async () => []);
            const cache = new RKeyCache(); // starts OFF (no servers)

            const before = await cache.resolveImageUrl(bridge as never, imageEl(), true);
            expect(before).toBe(NT_IMAGE_URL); // bare — fallback off
            expect(fetchSpy).not.toHaveBeenCalled();

            cache.setFallbackServers([SERVER]); // hot-reload pushes a new endpoint
            vi.setSystemTime(new Date(1_700_000_000_000 + 11_000)); // clear the 10s empty-cooldown

            const after = await cache.resolveImageUrl(bridge as never, imageEl(), true);
            expect(after).toBe(`${NT_IMAGE_URL}&rkey=HOTGROUP`);
            expect(fetchSpy).toHaveBeenCalledWith(SERVER, expect.anything());
        } finally {
            vi.useRealTimers();
        }
    });
});