proxy-station/server/src/routes/wireguard.test.ts

import { afterAll, beforeEach, describe, expect, it } from 'vitest';
import { mkdtempSync, rmSync } from 'node:fs';
import { tmpdir } from 'node:os';
import path from 'node:path';
import Database from 'better-sqlite3';
import { wireguardProfileSchema } from '@proxy-station/shared';

const directory = mkdtempSync(path.join(tmpdir(), 'proxy-station-wireguard-'));
process.env.DATA_DIR = directory;
const { db, createDb } = await import('../db/index.js');
const { nodes, subTokens, wgDevices, wgJobs, wgObservations, wgEvents, wgSubscriptionBindings, wgGateways, wgNetworks } = await import('../db/schema.js');
const { tokensRoute } = await import('./tokens.js');
const { subRoute } = await import('./sub.js');
const { generateSurge } = await import('../services/generators/surge.js');
const { generateShadowrocketConf } = await import('../services/generators/shadowrocket-conf.js');

const key = (value: number) => Buffer.alloc(32, value).toString('base64');
const profile = (address = '10.203.0.2', value = 2) => ({
  server: 'vpn.example.com', port: 51820, serverPublicKey: key(1), privateKey: key(value),
  address, subnet: '10.203.0.0/24', acceptIncoming: false,
});
const request = (method: string, body: unknown) => ({ method, headers: { 'content-type': 'application/json' }, body: JSON.stringify(body) });
beforeEach(() => {
  for (const table of [wgSubscriptionBindings, subTokens, wgEvents, wgObservations, wgJobs, wgDevices, wgGateways, wgNetworks, nodes]) db.delete(table).run();
});
afterAll(() => { db.$client.close(); rmSync(directory, { recursive: true, force: true }); });

describe('per-device WireGuard subscriptions', () => {
  it('isolates keys between device tokens and does not enable WireGuard on shared subscriptions', async () => {
    for (const [token, wireguard] of [['ipad-token', profile()], ['mac-token', profile('10.203.0.3', 3)], ['shared-token', null]] as const) {
      db.insert(subTokens).values({ id: token, token, name: token, createdAt: 0,
        wireguard: wireguard ? JSON.stringify(wireguard) : null }).run();
    }
    const ipad = await subRoute.request('/ipad-token/wireguard.conf');
    expect(ipad.headers.get('cache-control')).toBe('private, no-store');
    const config = await ipad.text();
    expect(config).toContain(`PrivateKey = ${key(2)}`);
    expect(config).not.toContain(key(3));
    expect(config).toContain('AllowedIPs = 10.203.0.0/24');
    expect(config).not.toContain('0.0.0.0/0');
    for (const token of ['shared-token', 'invalid-token']) {
      expect((await subRoute.request(`/${token}/wireguard.conf`)).status).toBe(404);
    }
    expect(await (await subRoute.request('/shared-token/surge')).text()).not.toContain('[WireGuard');
    for (const token of ['ipad-token', 'shared-token', 'invalid-token']) {
      expect((await subRoute.request(`/${token}/wireguard-shadowrocket`)).status).toBe(404);
    }
    db.update(subTokens).set({ enabled: false }).run();
    expect((await subRoute.request('/ipad-token/wireguard.conf')).status).toBe(404);
  });

  it('keeps Shadowrocket proxy downloads independent of legacy and paused WG identities', async () => {
    for (const [token, wireguard] of [['device', { ...profile(), presharedKey: key(9) }], ['shared', null]] as const) {
      db.insert(subTokens).values({ id: token, token, name: token, createdAt: 0,
        wireguard: wireguard ? JSON.stringify(wireguard) : null }).run();
    }
    db.insert(nodes).values({ id: 'proxy', name: 'Example Proxy Direct', protocol: 'trojan', server: 'proxy.example.test',
      port: 443, password: 'placeholder-password', createdAt: 0, updatedAt: 0 }).run();
    const plainNodes = await (await subRoute.request('/shared/shadowrocket')).text();
    expect(Buffer.from(plainNodes, 'base64').toString()).toContain('trojan://');
    const plainConf = await (await subRoute.request('/shared/shadowrocket-conf')).text();
    // A paused device must not gate the ordinary proxy subscription or trigger private-key retrieval.
    db.update(wgDevices).set({ desired: 'paused' }).run();
    for (const [endpoint, expected] of [['shadowrocket', plainNodes], ['shadowrocket-conf', plainConf]]) {
      for (const query of ['', '?wg=not-a-device-credential']) {
        const response = await subRoute.request(`/device/${endpoint}${query}`);
        expect(response.status).toBe(200);
        const content = await response.text();
        expect(content).toBe(expected);
        expect(content).not.toContain('WireGuard');
        expect(content).not.toContain('10.203.0.');
        expect(content).not.toContain(key(2));
        expect(content).not.toContain(key(9));
      }
    }
    const tokens = await (await tokensRoute.request('/')).json() as any[];
    expect(tokens.every(t => t.urls.shadowrocketWireguard === undefined)).toBe(true);
    expect(tokens.find(t => t.id === 'device').urls.wireguard).toBe('/sub/device/wireguard.conf');
  });

  it('prevents token edits from switching shared identity and preserves a legacy device on rename', async () => {
    db.insert(subTokens).values({ id: 'ipad', token: 'ipad-secret', name: 'iPad', createdAt: 0, wireguard: JSON.stringify(profile()) }).run();
    const created = await tokensRoute.request('/', request('POST', { name: 'shared' }));
    expect(created.status).toBe(201);
    const shared = await created.json() as any;
    expect((await tokensRoute.request('/', request('POST', { name: 'duplicate', wireguard: profile() }))).status).toBe(400);
    expect((await tokensRoute.request(`/${shared.id}`, request('PUT', { name: 'shared', wireguard: profile() }))).status).toBe(400);
    const updated = await tokensRoute.request('/ipad', request('PUT', { name: 'iPad renamed', nodeIds: [] }));
    const metadata = await updated.json() as any;
    expect(metadata.wgDeviceId).toBeTruthy();
    expect(metadata.wireguard).toBeUndefined();
    expect(await (await subRoute.request('/ipad-secret/wireguard.conf')).text()).toContain(key(2));
  });

  it('rejects default/public routes, invalid keys, injection and addresses outside the subnet', () => {
    for (const patch of [{ subnet: '0.0.0.0/0' }, { subnet: '8.8.8.0/24' }, { address: '10.204.0.2' },
      { address: '10.203.0.0' }, { privateKey: 'invalid' }, { server: 'vpn.example.com\n[Rule]' }]) {
      expect(wireguardProfileSchema.safeParse({ ...profile(), ...patch }).success).toBe(false);
    }
  });

  it('routes only the LAN before broad DIRECT rules, protects Internet groups and marks incoming requirements', () => {
    const ir = { nodes: [], sections: { General: '' },
      groups: [{ name: 'All', type: 'select' as const, members: [], includeAllNodes: true,
        filterRegex: null, testUrl: null, interval: null, tolerance: null }],
      rules: [{ type: 'IP-CIDR', value: '10.0.0.0/8', policy: 'DIRECT', params: [], ruleset: null },
        { type: 'FINAL', value: null, policy: 'DIRECT', params: [], ruleset: null }],
      wireguard: wireguardProfileSchema.parse({ ...profile(), acceptIncoming: true }),
    };
    const surge = generateSurge(ir, { selfUrl: 'https://example.com/sub/placeholder/surge' }).content;
    expect(surge.indexOf('IP-CIDR,10.203.0.0/24,NJCQ WireGuard LAN,no-resolve')).toBeLessThan(surge.indexOf('IP-CIDR,10.0.0.0/8'));
    expect(surge).toContain('accept-incoming = true');
    expect(surge).toContain('Surge Mac 6.11.0+');
    expect(surge.split('[Proxy Group]')[1].split('[Rule]')[0]).not.toContain('NJCQ WireGuard LAN');
    const dynamic = generateSurge({ ...ir, groups: [{ ...ir.groups[0], filterRegex: 'US|LAN' }] }, { selfUrl: 'https://example.test/sub/placeholder/surge' }).content;
    const filter = dynamic.match(/policy-regex-filter=([^\n]+)/)![1];
    const decodedFilter = filter.startsWith('"') ? JSON.parse(filter) : filter;
    expect(new RegExp(decodedFilter).test('NJCQ WireGuard LAN')).toBe(false);
    expect(new RegExp(decodedFilter).test('US Test Direct')).toBe(true);
    expect(new RegExp(decodedFilter).test('prefix US Test Direct')).toBe(true);
    expect(new RegExp(decodedFilter).test('prefix LAN')).toBe(true);
    expect(new RegExp(decodedFilter).test('JP Direct')).toBe(false);
    expect(generateShadowrocketConf(ir)).toEqual(generateShadowrocketConf({ ...ir, wireguard: undefined }));
  });

  it('upgrades an existing database without changing shared token identity', () => {
    const file = path.join(directory, 'legacy.db');
    const old = new Database(file);
    old.exec("CREATE TABLE sub_tokens (id TEXT PRIMARY KEY, token TEXT NOT NULL UNIQUE, name TEXT NOT NULL, enabled INTEGER NOT NULL DEFAULT 1, created_at INTEGER NOT NULL, last_access_at INTEGER, access_count INTEGER NOT NULL DEFAULT 0); INSERT INTO sub_tokens(id,token,name,created_at) VALUES ('legacy','old-token','existing',0)");
    old.close();
    const migrated = createDb(file);
    const token = migrated.select().from(subTokens).get()!;
    expect(token.token).toBe('old-token');
    expect(token.wireguard).toBeNull();
    migrated.$client.close();
  });
});