proxy-station/server/src/services/wg.ts

import { and, asc, desc, eq, isNull } from 'drizzle-orm';
import { nanoid } from 'nanoid';
import { wireguardProfileSchema, type WireguardProfile, type WgDeviceDto } from '@proxy-station/shared';
import { db } from '../db/index.js';
import { subTokens, wgDevices, wgEvents, wgGateways, wgJobs, wgNetworks, wgObservations, wgSubscriptionBindings } from '../db/schema.js';
import { credential, decryptSecret, encryptSecret, hashCredential, newIdentity, publicKey } from './wg-crypto.js';
import { migrateWireguard } from './wg-migrate.js';

export class WgError extends Error { constructor(message: string, public status: 404 | 409 | 503 = 409) { super(message); } }
export const migrateLegacy = () => migrateWireguard(db);
export function getDevice(id: string) {
  const row = db.select().from(wgDevices).where(eq(wgDevices.id, id)).get();
  if (!row) throw new WgError('设备不存在', 404);
  return row;
}
export function profileFor(device: typeof wgDevices.$inferSelect) {
  return wireguardProfileSchema.parse(decryptSecret<WireguardProfile>(device.secret));
}
export function canDistribute(device: typeof wgDevices.$inferSelect, legacy = false) {
  const job = db.select().from(wgJobs).where(eq(wgJobs.deviceId, device.id)).orderBy(desc(wgJobs.version)).get();
  return device.desired === 'active' && (device.appliedVersion === device.version ||
    (legacy && device.version === 1 && job?.action === 'adopt' && job.state === 'pending'));
}
export function authorizedProfile(id: string, legacy = false) {
  const device = getDevice(id);
  if (!canDistribute(device, legacy)) throw new WgError('设备尚未生效、已暂停或已移除', 404);
  const profile = profileFor(device);
  // New device delivery always enables Surge inbound; legacy URLs retain their original setting.
  return legacy ? profile : { ...profile, acceptIncoming: true };
}
export function deviceDto(device: typeof wgDevices.$inferSelect, now = Date.now()): WgDeviceDto {
  const observation = db.select().from(wgObservations).where(eq(wgObservations.deviceId, device.id)).get();
  const job = db.select().from(wgJobs).where(eq(wgJobs.deviceId, device.id)).orderBy(desc(wgJobs.version)).get();
  const fresh = !!observation && now - observation.receivedAt <= 45_000;
  let status: WgDeviceDto['status'];
  if (device.appliedVersion !== device.version) status = job?.state === 'failed' ? 'error' :
    device.desired === 'active' ? (job?.action === 'resume' ? 'restoring' : 'provisioning') : 'revoking';
  else if (device.desired !== 'active') status = device.desired;
  else if (!fresh) status = 'unknown';
  else if (!observation.handshakeAt && !observation.activityAt) status = 'pending';
  else status = now - Math.max(observation.handshakeAt ?? 0, observation.activityAt ?? 0) <= 180_000 ? 'online' : 'offline';
  return { id: device.id, gatewayId: device.gatewayId, name: device.name, os: device.os, address: device.address,
    publicKey: device.publicKey, status, canDistribute: canDistribute(device), createdAt: device.createdAt,
    sampledAt: observation?.receivedAt ?? null, lastActivityAt: observation?.activityAt ?? observation?.handshakeAt ?? null,
    uploadRate: fresh && status !== 'paused' ? observation.uploadRate : null,
    downloadRate: fresh && status !== 'paused' ? observation.downloadRate : null,
    uploadTotal: observation?.uploadTotal ?? 0, downloadTotal: observation?.downloadTotal ?? 0,
    job: job ? { id: job.id, action: job.action, state: job.state, error: job.error } : null };
}
export function listDevices() { migrateLegacy(); return db.select().from(wgDevices).where(isNull(wgDevices.archivedAt)).orderBy(asc(wgDevices.createdAt)).all().map(d => deviceDto(d)); }
function event(id: string, message: string) { db.insert(wgEvents).values({ id: nanoid(), deviceId: id, message, createdAt: Date.now() }).run(); }
export function createDevice(input: { name: string; os: string; requestId: string; gatewayId?: string }) {
  migrateLegacy();
  return db.transaction(() => {
    // Stable client request ID survives network timeout/double submission.
    const prior = db.select().from(wgDevices).where(eq(wgDevices.id, input.requestId)).get();
    if (prior) return deviceDto(prior);
    if (db.select().from(wgDevices).where(and(eq(wgDevices.name, input.name), isNull(wgDevices.archivedAt))).get()) throw new WgError('已有同名设备，请用位置或用途区分');
    const gateways = db.select().from(wgGateways).all();
    const gateway = input.gatewayId ? gateways.find(g => g.id === input.gatewayId) : gateways.length === 1 ? gateways[0] : null;
    if (!gateway?.agentHash) throw new WgError('请先配置网关及独立 agent 凭证', 503);
    const network = db.select().from(wgNetworks).where(eq(wgNetworks.id, gateway.networkId)).get()!;
    if (network.subnet !== '10.203.0.0/24') throw new WgError('当前仅开放 njcq 设备网络');
    const used = new Set(db.select().from(wgDevices).all().map(d => d.address));
    // .2/.3 reserved for the existing iPad/Mac even before adoption; never recycle archives.
    const address = Array.from({ length: 250 }, (_, i) => `10.203.0.${i + 4}`).find(ip => !used.has(ip));
    if (!address) throw new WgError('设备地址已用尽');
    const profile = wireguardProfileSchema.parse({ server: gateway.server, port: gateway.port, serverPublicKey: gateway.publicKey,
      ...newIdentity(), address, subnet: network.subnet, mtu: network.mtu, acceptIncoming: true });
    const device = { id: input.requestId, gatewayId: gateway.id, name: input.name, os: input.os, address,
      publicKey: publicKey(profile.privateKey), secret: encryptSecret(profile), desired: 'active' as const,
      version: 1, appliedVersion: 0, createdAt: Date.now(), archivedAt: null };
    db.insert(wgDevices).values(device).run();
    db.insert(wgJobs).values({ id: nanoid(), deviceId: device.id, action: 'create', version: 1, state: 'pending', createdAt: Date.now() }).run();
    event(device.id, '已预留固定 IP 和身份，等待网关确认');
    return deviceDto(device);
  });
}
export function deviceAction(id: string, action: 'pause' | 'resume' | 'rotate' | 'remove' | 'retry') {
  return db.transaction(() => {
    const device = getDevice(id);
    const previous = db.select().from(wgJobs).where(eq(wgJobs.deviceId, id)).orderBy(desc(wgJobs.version)).get();
    if (device.archivedAt) throw new WgError('设备已移除');
    if (device.appliedVersion !== device.version) {
      if (previous?.action === action) return deviceDto(device);
      if (action !== 'retry') throw new WgError('请先完成或重试当前网关任务');
      db.update(wgJobs).set({ state: 'pending', error: null }).where(eq(wgJobs.id, previous!.id)).run();
      return deviceDto(device);
    }
    if (action === 'retry') return deviceDto(device);
    if (action === 'resume' && device.desired !== 'paused') throw new WgError('设备未暂停');
    if (action === 'pause' && device.desired === 'paused') return deviceDto(device);
    if (action === 'rotate' && device.desired !== 'active') throw new WgError('请先恢复设备');
    const desired = action === 'pause' ? 'paused' : action === 'remove' ? 'archived' : 'active';
    const patch = action === 'rotate' ? (() => {
      const profile = { ...profileFor(device), ...newIdentity() };
      return { publicKey: publicKey(profile.privateKey), secret: encryptSecret(profile) };
    })() : {};
    db.update(wgDevices).set({ ...patch, desired, version: device.version + 1 }).where(eq(wgDevices.id, id)).run();
    db.insert(wgJobs).values({ id: nanoid(), deviceId: id, action, version: device.version + 1, state: 'pending',
      oldPublicKey: action === 'rotate' ? device.publicKey : null, createdAt: Date.now() }).run();
    event(id, `已请求 ${action}，等待运行与持久配置确认`);
    return deviceDto(getDevice(id));
  });
}
export function issueBinding(tokenId: string, deviceId: string, purpose: 'surge' | 'clash' = 'surge') {
  const token = db.select().from(subTokens).where(eq(subTokens.id, tokenId)).get();
  if (!token?.enabled) throw new WgError('订阅不存在或已停用', 404);
  authorizedProfile(deviceId);
  const secret = credential(); const id = nanoid();
  db.insert(wgSubscriptionBindings).values({ id, tokenId, deviceId, credentialHash: hashCredential(secret),
    purpose, createdAt: Date.now() }).run();
  return { id, deviceId, path: `/sub/${token.token}/${purpose}?wg=${secret}` };
}
export function bindingProfile(tokenId: string, secret: string, purpose: 'surge' | 'clash' = 'surge') {
  if (!/^[\w-]{43}$/.test(secret)) throw new WgError('设备凭证无效', 404);
  const binding = db.select().from(wgSubscriptionBindings).where(and(eq(wgSubscriptionBindings.credentialHash, hashCredential(secret)),
    eq(wgSubscriptionBindings.tokenId, tokenId), eq(wgSubscriptionBindings.purpose, purpose), isNull(wgSubscriptionBindings.revokedAt))).get();
  if (!binding) throw new WgError('设备凭证无效', 404);
  const profile = authorizedProfile(binding.deviceId);
  db.update(wgSubscriptionBindings).set({ lastAccessAt: Date.now() }).where(eq(wgSubscriptionBindings.id, binding.id)).run();
  return profile;
}

// Agent requests contain public keys/PSKs but never client private keys.
export function gatewayTasks(gatewayId: string) {
  const jobs = db.select().from(wgJobs).innerJoin(wgDevices, eq(wgJobs.deviceId, wgDevices.id))
    .where(and(eq(wgDevices.gatewayId, gatewayId), eq(wgJobs.state, 'pending'))).all();
  return jobs.map(({ wg_jobs: j, wg_devices: d }) => ({ id: j.id, deviceId: d.id, action: j.action,
    version: j.version, publicKey: d.publicKey, oldPublicKey: j.oldPublicKey,
    address: d.address, presharedKey: profileFor(d).presharedKey ?? null, desired: d.desired }));
}
export interface AgentReport {
  bootId: string;
  receipts: { id: string; version: number; runtime: boolean; persistent: boolean; error?: string }[];
  peers: { publicKey: string; handshakeAge: number | null; rx: number; tx: number }[];
}
export function receiveReport(gatewayId: string, report: AgentReport) {
  const now = Date.now();
  db.transaction(() => {
    db.update(wgGateways).set({ receivedAt: now, bootId: report.bootId }).where(eq(wgGateways.id, gatewayId)).run();
    for (const receipt of report.receipts) {
      const job = db.select().from(wgJobs).where(eq(wgJobs.id, receipt.id)).get();
      if (!job || job.state === 'confirmed' || job.version !== receipt.version) continue;
      const device = getDevice(job.deviceId);
      if (device.gatewayId !== gatewayId || device.version !== receipt.version) continue;
      const ok = receipt.runtime && receipt.persistent && !receipt.error;
      db.update(wgJobs).set({ state: ok ? 'confirmed' : 'failed', error: ok ? null : '网关配置未确认，请重试任务',
        confirmedAt: ok ? now : null }).where(eq(wgJobs.id, job.id)).run();
      if (ok) {
        db.update(wgDevices).set({ appliedVersion: job.version, ...(device.desired === 'archived' ? { archivedAt: now } : {}) })
          .where(eq(wgDevices.id, device.id)).run();
        if (device.desired === 'archived') db.update(wgSubscriptionBindings).set({ revokedAt: now }).where(eq(wgSubscriptionBindings.deviceId, device.id)).run();
        if (job.action === 'rotate') db.delete(wgObservations).where(eq(wgObservations.deviceId, device.id)).run();
      }
      event(device.id, ok ? '网关确认：运行配置与持久配置均已生效' : '网关未确认配置，保留当前身份与任务供重试');
    }
    const peers = new Map(report.peers.map(p => [p.publicKey, p]));
    for (const device of db.select().from(wgDevices).where(eq(wgDevices.gatewayId, gatewayId)).all()) {
      const peer = peers.get(device.publicKey);
      if (!peer || device.desired !== 'active') continue;
      const old = db.select().from(wgObservations).where(eq(wgObservations.deviceId, device.id)).get();
      const elapsed = old ? (now - old.receivedAt) / 1000 : 0;
      const sameEpoch = old?.bootId === report.bootId && peer.rx >= old.rx && peer.tx >= old.tx;
      const up = sameEpoch ? peer.rx - old!.rx : 0; // gateway rx = device upload
      const down = sameEpoch ? peer.tx - old!.tx : 0;
      const fresh = elapsed > 0 && elapsed <= 45;
      const handshakeAt = peer.handshakeAge === null ? old?.handshakeAt ?? null : now - peer.handshakeAge * 1000;
      const activityAt = (up || down) ? now : Math.max(old?.activityAt ?? 0, handshakeAt ?? 0) || null;
      const row = { deviceId: device.id, receivedAt: now, bootId: report.bootId, handshakeAt, activityAt,
        rx: peer.rx, tx: peer.tx, uploadRate: fresh && sameEpoch ? Math.round(up / elapsed) : null,
        downloadRate: fresh && sameEpoch ? Math.round(down / elapsed) : null,
        uploadTotal: (old?.uploadTotal ?? 0) + up, downloadTotal: (old?.downloadTotal ?? 0) + down };
      db.insert(wgObservations).values(row).onConflictDoUpdate({ target: wgObservations.deviceId, set: row }).run();
    }
  });
}