dcrouter/ts/vpn/classes.vpn-manager.ts

import * as plugins from '../plugins.js';
import { logger } from '../logger.js';
import { VpnServerKeysDoc, VpnClientDoc } from '../db/index.js';

export interface IVpnManagerConfig {
  /** VPN subnet CIDR (default: '10.8.0.0/24') */
  subnet?: string;
  /** WireGuard UDP listen port (default: 51820) */
  wgListenPort?: number;
  /** DNS servers pushed to VPN clients */
  dns?: string[];
  /** Server endpoint hostname for client configs (e.g. 'vpn.example.com') */
  serverEndpoint?: string;
  /** Pre-defined VPN clients created on startup (idempotent — skips already-persisted clients) */
  initialClients?: Array<{
    clientId: string;
    targetProfileIds?: string[];
    description?: string;
  }>;
  /** Called when clients are created/deleted/toggled — triggers route re-application */
  onClientChanged?: () => void;
  /** Destination routing policy override. Default: forceTarget to 127.0.0.1 */
  destinationPolicy?: {
    default: 'forceTarget' | 'block' | 'allow';
    target?: string;
    allowList?: string[];
    blockList?: string[];
  };
  /** Compute per-client AllowedIPs based on the client's target profile IDs.
   *  Called at config generation time (create/export). Returns CIDRs for WireGuard AllowedIPs.
   *  When not set, defaults to [subnet]. */
  getClientAllowedIPs?: (targetProfileIds: string[]) => Promise<string[]>;
  /** Resolve per-client destination allow-list IPs from target profile IDs.
   *  Returns IP strings that should bypass forceTarget and go direct to the real destination. */
  getClientDirectTargets?: (targetProfileIds: string[]) => string[];
  /** Forwarding mode: 'socket' (default, userspace NAT), 'bridge' (L2 bridge to host LAN),
   *  or 'hybrid' (socket default, bridge for clients with useHostIp=true) */
  forwardingMode?: 'socket' | 'bridge' | 'hybrid';
  /** LAN subnet CIDR for bridge mode (e.g., '192.168.1.0/24') */
  bridgeLanSubnet?: string;
  /** Physical network interface for bridge mode (auto-detected if omitted) */
  bridgePhysicalInterface?: string;
  /** Start of VPN client IP range in LAN subnet (host offset, default: 200) */
  bridgeIpRangeStart?: number;
  /** End of VPN client IP range in LAN subnet (host offset, default: 250) */
  bridgeIpRangeEnd?: number;
}

/**
 * Manages the SmartVPN server lifecycle and VPN client CRUD.
 * Persists server keys and client registrations via smartdata document classes.
 */
export class VpnManager {
  private config: IVpnManagerConfig;
  private vpnServer?: plugins.smartvpn.VpnServer;
  private clients: Map<string, VpnClientDoc> = new Map();
  private serverKeys?: VpnServerKeysDoc;
  private resolvedForwardingMode?: 'socket' | 'bridge' | 'hybrid';
  private forwardingModeOverride?: 'socket' | 'bridge' | 'hybrid';

  constructor(config: IVpnManagerConfig) {
    this.config = config;
  }

  /** The VPN subnet CIDR. */
  public getSubnet(): string {
    return this.config.subnet || '10.8.0.0/24';
  }

  /** Whether the VPN server is running. */
  public get running(): boolean {
    return this.vpnServer?.running ?? false;
  }

  /**
   * Start the VPN server.
   * Loads or generates server keys, loads persisted clients, starts VpnServer.
   */
  public async start(): Promise<void> {
    // Load or generate server keys
    this.serverKeys = await this.loadOrGenerateServerKeys();

    // Load persisted clients
    await this.loadPersistedClients();

    // Build client entries for the daemon
    const clientEntries: plugins.smartvpn.IClientEntry[] = [];
    let anyClientUsesHostIp = false;
    for (const client of this.clients.values()) {
      if (client.useHostIp) {
        anyClientUsesHostIp = true;
      }
      this.normalizeClientRoutingSettings(client);
      const entry: plugins.smartvpn.IClientEntry = {
        clientId: client.clientId,
        publicKey: client.noisePublicKey,
        wgPublicKey: client.wgPublicKey,
        enabled: client.enabled,
        description: client.description,
        assignedIp: client.assignedIp,
        expiresAt: client.expiresAt,
        security: this.buildClientSecurity(client),
        useHostIp: client.useHostIp,
        useDhcp: client.useDhcp,
        staticIp: client.staticIp,
        forceVlan: client.forceVlan,
        vlanId: client.vlanId,
      };
      clientEntries.push(entry);
    }

    const subnet = this.getSubnet();
    const wgListenPort = this.config.wgListenPort ?? 51820;

    // Auto-detect hybrid mode: if any persisted client uses host IP and mode is
    // 'socket' (or unset), upgrade to 'hybrid' so the daemon can handle both
    let configuredMode = this.forwardingModeOverride ?? this.config.forwardingMode ?? 'socket';
    if (anyClientUsesHostIp && configuredMode === 'socket') {
      configuredMode = 'hybrid';
      logger.log('info', 'VPN: Auto-upgrading forwarding mode to hybrid (client with useHostIp detected)');
    }
    const forwardingMode = configuredMode === 'hybrid' ? 'hybrid' : configuredMode;
    const isBridge = forwardingMode === 'bridge';
    this.resolvedForwardingMode = forwardingMode;
    this.forwardingModeOverride = undefined;

    // Create and start VpnServer
    this.vpnServer = new plugins.smartvpn.VpnServer({
      transport: { transport: 'stdio' },
    });

    // Default destination policy: bridge mode allows traffic through directly,
    // socket mode forces traffic to SmartProxy on 127.0.0.1
    const defaultDestinationPolicy: plugins.smartvpn.IDestinationPolicy = isBridge
      ? { default: 'allow' as const }
      : { default: 'forceTarget' as const, target: '127.0.0.1' };

    const serverConfig: plugins.smartvpn.IVpnServerConfig = {
      listenAddr: '0.0.0.0:0', // WS listener not strictly needed but required field
      privateKey: this.serverKeys.noisePrivateKey,
      publicKey: this.serverKeys.noisePublicKey,
      subnet,
      dns: this.config.dns,
      forwardingMode: forwardingMode as any,
      transportMode: 'all',
      wgPrivateKey: this.serverKeys.wgPrivateKey,
      wgListenPort,
      clients: clientEntries,
      socketForwardProxyProtocol: !isBridge,
      destinationPolicy: this.getServerDestinationPolicy(forwardingMode, defaultDestinationPolicy),
      serverEndpoint: this.config.serverEndpoint
        ? `${this.config.serverEndpoint}:${wgListenPort}`
        : undefined,
      clientAllowedIPs: [subnet],
      // Bridge-specific config
      ...(isBridge ? {
        bridgeLanSubnet: this.config.bridgeLanSubnet,
        bridgePhysicalInterface: this.config.bridgePhysicalInterface,
        bridgeIpRangeStart: this.config.bridgeIpRangeStart,
        bridgeIpRangeEnd: this.config.bridgeIpRangeEnd,
      } : {}),
    };

    await this.vpnServer.start(serverConfig);

    // Create initial clients from config (idempotent — skip already-persisted)
    if (this.config.initialClients) {
      for (const initial of this.config.initialClients) {
        if (!this.clients.has(initial.clientId)) {
          const bundle = await this.createClient({
            clientId: initial.clientId,
            targetProfileIds: initial.targetProfileIds,
            description: initial.description,
          });
          logger.log('info', `VPN: Created initial client '${initial.clientId}' (IP: ${bundle.entry.assignedIp})`);
        }
      }
    }

    logger.log('info', `VPN server started: subnet=${subnet}, wg=:${wgListenPort}, clients=${this.clients.size}`);
  }

  /**
   * Stop the VPN server.
   */
  public async stop(): Promise<void> {
    if (this.vpnServer) {
      try {
        await this.vpnServer.stopServer();
      } catch {
        // Ignore stop errors
      }
      this.vpnServer.stop();
      this.vpnServer = undefined;
    }
    this.resolvedForwardingMode = undefined;
    logger.log('info', 'VPN server stopped');
  }

  // ── Client CRUD ────────────────────────────────────────────────────────

  /**
   * Create a new VPN client. Returns the config bundle (secrets only shown once).
   */
  public async createClient(opts: {
    clientId: string;
    targetProfileIds?: string[];
    description?: string;
    destinationAllowList?: string[];
    destinationBlockList?: string[];
    useHostIp?: boolean;
    useDhcp?: boolean;
    staticIp?: string;
    forceVlan?: boolean;
    vlanId?: number;
  }): Promise<plugins.smartvpn.IClientConfigBundle> {
    if (!this.vpnServer) {
      throw new Error('VPN server not running');
    }

    await this.ensureForwardingModeForHostIpClient(opts.useHostIp === true);

    const doc = new VpnClientDoc();
    doc.clientId = opts.clientId;
    doc.enabled = true;
    doc.targetProfileIds = opts.targetProfileIds;
    doc.description = opts.description;
    doc.destinationAllowList = opts.destinationAllowList;
    doc.destinationBlockList = opts.destinationBlockList;
    doc.useHostIp = opts.useHostIp;
    doc.useDhcp = opts.useDhcp;
    doc.staticIp = opts.staticIp;
    doc.forceVlan = opts.forceVlan;
    doc.vlanId = opts.vlanId;
    doc.createdAt = Date.now();
    doc.updatedAt = Date.now();
    this.normalizeClientRoutingSettings(doc);

    const bundle = await this.vpnServer.createClient({
      clientId: doc.clientId,
      description: doc.description,
      security: this.buildClientSecurity(doc),
      useHostIp: doc.useHostIp,
      useDhcp: doc.useDhcp,
      staticIp: doc.staticIp,
      forceVlan: doc.forceVlan,
      vlanId: doc.vlanId,
    });

    // Override AllowedIPs with per-client values based on target profiles
    if (this.config.getClientAllowedIPs && bundle.wireguardConfig) {
      const allowedIPs = await this.config.getClientAllowedIPs(doc.targetProfileIds || []);
      bundle.wireguardConfig = bundle.wireguardConfig.replace(
        /AllowedIPs\s*=\s*.+/,
        `AllowedIPs = ${allowedIPs.join(', ')}`,
      );
    }

    // Persist client entry (including WG private key for export/QR)
    doc.clientId = bundle.entry.clientId;
    doc.enabled = bundle.entry.enabled ?? true;
    doc.description = bundle.entry.description;
    doc.assignedIp = bundle.entry.assignedIp;
    doc.noisePublicKey = bundle.entry.publicKey;
    doc.wgPublicKey = bundle.entry.wgPublicKey || '';
    doc.wgPrivateKey = bundle.secrets?.wgPrivateKey
      || bundle.wireguardConfig?.match(/PrivateKey\s*=\s*(.+)/)?.[1]?.trim();
    doc.updatedAt = Date.now();
    doc.expiresAt = bundle.entry.expiresAt;
    this.clients.set(doc.clientId, doc);
    try {
      await this.persistClient(doc);
    } catch (err) {
      // Rollback: remove from in-memory map and daemon to stay consistent with DB
      this.clients.delete(doc.clientId);
      try {
        await this.vpnServer!.removeClient(doc.clientId);
      } catch {
        // best-effort daemon cleanup
      }
      throw err;
    }

    this.config.onClientChanged?.();
    return bundle;
  }

  /**
   * Remove a VPN client.
   */
  public async removeClient(clientId: string): Promise<void> {
    if (!this.vpnServer) {
      throw new Error('VPN server not running');
    }
    await this.vpnServer.removeClient(clientId);
    const doc = this.clients.get(clientId);
    this.clients.delete(clientId);
    if (doc) {
      await doc.delete();
    }
    this.config.onClientChanged?.();
  }

  /**
   * List all registered clients (without secrets).
   */
  public listClients(): VpnClientDoc[] {
    return [...this.clients.values()];
  }

  /**
   * Enable a client.
   */
  public async enableClient(clientId: string): Promise<void> {
    if (!this.vpnServer) throw new Error('VPN server not running');
    await this.vpnServer.enableClient(clientId);
    const client = this.clients.get(clientId);
    if (client) {
      client.enabled = true;
      client.updatedAt = Date.now();
      await this.persistClient(client);
    }
    this.config.onClientChanged?.();
  }

  /**
   * Disable a client.
   */
  public async disableClient(clientId: string): Promise<void> {
    if (!this.vpnServer) throw new Error('VPN server not running');
    await this.vpnServer.disableClient(clientId);
    const client = this.clients.get(clientId);
    if (client) {
      client.enabled = false;
      client.updatedAt = Date.now();
      await this.persistClient(client);
    }
    this.config.onClientChanged?.();
  }

  /**
   * Update a client's metadata (description, target profiles) without rotating keys.
   */
  public async updateClient(clientId: string, update: {
    description?: string;
    targetProfileIds?: string[];
    destinationAllowList?: string[];
    destinationBlockList?: string[];
    useHostIp?: boolean;
    useDhcp?: boolean;
    staticIp?: string;
    forceVlan?: boolean;
    vlanId?: number;
  }): Promise<void> {
    const client = this.clients.get(clientId);
    if (!client) throw new Error(`Client not found: ${clientId}`);
    if (update.description !== undefined) client.description = update.description;
    if (update.targetProfileIds !== undefined) client.targetProfileIds = update.targetProfileIds;
    if (update.destinationAllowList !== undefined) client.destinationAllowList = update.destinationAllowList;
    if (update.destinationBlockList !== undefined) client.destinationBlockList = update.destinationBlockList;
    if (update.useHostIp !== undefined) client.useHostIp = update.useHostIp;
    if (update.useDhcp !== undefined) client.useDhcp = update.useDhcp;
    if (update.staticIp !== undefined) client.staticIp = update.staticIp;
    if (update.forceVlan !== undefined) client.forceVlan = update.forceVlan;
    if (update.vlanId !== undefined) client.vlanId = update.vlanId;
    this.normalizeClientRoutingSettings(client);
    client.updatedAt = Date.now();
    await this.persistClient(client);

    if (this.vpnServer) {
      await this.ensureForwardingModeForHostIpClient(client.useHostIp === true);
      await this.vpnServer.updateClient(clientId, this.buildClientRuntimeUpdate(client));
    }

    this.config.onClientChanged?.();
  }

  /**
   * Rotate a client's keys. Returns the new config bundle.
   */
  public async rotateClientKey(clientId: string): Promise<plugins.smartvpn.IClientConfigBundle> {
    if (!this.vpnServer) throw new Error('VPN server not running');
    const bundle = await this.vpnServer.rotateClientKey(clientId);

    // Update persisted entry with new keys (including private key for export/QR)
    const client = this.clients.get(clientId);
    if (client) {
      client.noisePublicKey = bundle.entry.publicKey;
      client.wgPublicKey = bundle.entry.wgPublicKey || '';
      client.wgPrivateKey = bundle.secrets?.wgPrivateKey
        || bundle.wireguardConfig?.match(/PrivateKey\s*=\s*(.+)/)?.[1]?.trim();
      client.updatedAt = Date.now();
      await this.persistClient(client);
    }

    return bundle;
  }

  /**
   * Export a client config. Injects stored WG private key and per-client AllowedIPs.
   */
  public async exportClientConfig(clientId: string, format: 'smartvpn' | 'wireguard'): Promise<string> {
    if (!this.vpnServer) throw new Error('VPN server not running');
    let config = await this.vpnServer.exportClientConfig(clientId, format);

    if (format === 'wireguard') {
      const persisted = this.clients.get(clientId);

      // Inject stored WG private key so exports produce valid, scannable configs
      if (persisted?.wgPrivateKey) {
        config = config.replace(
          '[Interface]\n',
          `[Interface]\nPrivateKey = ${persisted.wgPrivateKey}\n`,
        );
      }

      // Override AllowedIPs with per-client values based on target profiles
      if (this.config.getClientAllowedIPs) {
        const profileIds = persisted?.targetProfileIds || [];
        const allowedIPs = await this.config.getClientAllowedIPs(profileIds);
        config = config.replace(
          /AllowedIPs\s*=\s*.+/,
          `AllowedIPs = ${allowedIPs.join(', ')}`,
        );
      }
    }

    return config;
  }

  // ── Status and telemetry ───────────────────────────────────────────────

  /**
   * Get server status.
   */
  public async getStatus(): Promise<plugins.smartvpn.IVpnStatus | null> {
    if (!this.vpnServer) return null;
    return this.vpnServer.getStatus();
  }

  /**
   * Get server statistics.
   */
  public async getStatistics(): Promise<plugins.smartvpn.IVpnServerStatistics | null> {
    if (!this.vpnServer) return null;
    return this.vpnServer.getStatistics();
  }

  /**
   * List currently connected clients.
   */
  public async getConnectedClients(): Promise<plugins.smartvpn.IVpnClientInfo[]> {
    if (!this.vpnServer) return [];
    return this.vpnServer.listClients();
  }

  /**
   * Get telemetry for a specific client.
   */
  public async getClientTelemetry(clientId: string): Promise<plugins.smartvpn.IVpnClientTelemetry | null> {
    if (!this.vpnServer) return null;
    return this.vpnServer.getClientTelemetry(clientId);
  }

  /**
   * Get server public keys (for display/info).
   */
  public getServerPublicKeys(): { noisePublicKey: string; wgPublicKey: string } | null {
    if (!this.serverKeys) return null;
    return {
      noisePublicKey: this.serverKeys.noisePublicKey,
      wgPublicKey: this.serverKeys.wgPublicKey,
    };
  }

  // ── Per-client security ────────────────────────────────────────────────

  /**
   * Build per-client security settings for the smartvpn daemon.
   * TargetProfile direct IP:port targets extend the effective allow-list.
   */
  private buildClientSecurity(client: VpnClientDoc): plugins.smartvpn.IClientSecurity {
    const security: plugins.smartvpn.IClientSecurity = {};
    const basePolicy = this.getBaseDestinationPolicy(client);

    const profileDirectTargets = this.config.getClientDirectTargets?.(client.targetProfileIds || []) || [];
    const mergedAllowList = this.mergeDestinationLists(
      basePolicy.allowList,
      client.destinationAllowList,
      profileDirectTargets,
    );
    const mergedBlockList = this.mergeDestinationLists(
      basePolicy.blockList,
      client.destinationBlockList,
    );

    security.destinationPolicy = {
      default: basePolicy.default,
      target: basePolicy.default === 'forceTarget' ? basePolicy.target : undefined,
      allowList: mergedAllowList.length ? mergedAllowList : undefined,
      blockList: mergedBlockList.length ? mergedBlockList : undefined,
    };

    return security;
  }

  /**
   * Refresh all client security policies against the running daemon.
   * Call this when TargetProfiles change so destination allow-lists stay in sync.
   */
  public async refreshAllClientSecurity(): Promise<void> {
    if (!this.vpnServer) return;
    for (const client of this.clients.values()) {
      await this.vpnServer.updateClient(client.clientId, this.buildClientRuntimeUpdate(client));
    }
  }

  // ── Private helpers ────────────────────────────────────────────────────

  private async loadOrGenerateServerKeys(): Promise<VpnServerKeysDoc> {
    const stored = await VpnServerKeysDoc.load();
    if (stored?.noisePrivateKey && stored?.wgPrivateKey) {
      logger.log('info', 'Loaded VPN server keys from storage');
      return stored;
    }

    // Generate new keys via the daemon
    const tempServer = new plugins.smartvpn.VpnServer({
      transport: { transport: 'stdio' },
    });
    await tempServer.start();

    const noiseKeys = await tempServer.generateKeypair();
    const wgKeys = await tempServer.generateWgKeypair();
    tempServer.stop();

    const doc = stored || new VpnServerKeysDoc();
    doc.noisePrivateKey = noiseKeys.privateKey;
    doc.noisePublicKey = noiseKeys.publicKey;
    doc.wgPrivateKey = wgKeys.privateKey;
    doc.wgPublicKey = wgKeys.publicKey;
    await doc.save();

    logger.log('info', 'Generated and persisted new VPN server keys');
    return doc;
  }

  private async loadPersistedClients(): Promise<void> {
    const docs = await VpnClientDoc.findAll();
    for (const doc of docs) {
      this.normalizeClientRoutingSettings(doc);
      this.clients.set(doc.clientId, doc);
    }
    if (this.clients.size > 0) {
      logger.log('info', `Loaded ${this.clients.size} persisted VPN client(s)`);
    }
  }

  private getResolvedForwardingMode(): 'socket' | 'bridge' | 'hybrid' {
    return this.resolvedForwardingMode
      ?? this.forwardingModeOverride
      ?? this.config.forwardingMode
      ?? 'socket';
  }

  private getDefaultDestinationPolicy(
    forwardingMode: 'socket' | 'bridge' | 'hybrid',
    useHostIp = false,
  ): plugins.smartvpn.IDestinationPolicy {
    if (forwardingMode === 'bridge' || (forwardingMode === 'hybrid' && useHostIp)) {
      return { default: 'allow' };
    }
    return { default: 'forceTarget', target: '127.0.0.1' };
  }

  private getServerDestinationPolicy(
    forwardingMode: 'socket' | 'bridge' | 'hybrid',
    fallbackPolicy = this.getDefaultDestinationPolicy(forwardingMode),
  ): plugins.smartvpn.IDestinationPolicy {
    return this.config.destinationPolicy ?? fallbackPolicy;
  }

  private getBaseDestinationPolicy(client: Pick<VpnClientDoc, 'useHostIp'>): plugins.smartvpn.IDestinationPolicy {
    if (this.config.destinationPolicy) {
      return { ...this.config.destinationPolicy };
    }
    return this.getDefaultDestinationPolicy(this.getResolvedForwardingMode(), client.useHostIp === true);
  }

  private mergeDestinationLists(...lists: Array<string[] | undefined>): string[] {
    const merged = new Set<string>();
    for (const list of lists) {
      for (const entry of list || []) {
        merged.add(entry);
      }
    }
    return [...merged];
  }

  private normalizeClientRoutingSettings(
    client: Pick<VpnClientDoc, 'useHostIp' | 'useDhcp' | 'staticIp' | 'forceVlan' | 'vlanId'>,
  ): void {
    client.useHostIp = client.useHostIp === true;

    if (!client.useHostIp) {
      client.useDhcp = false;
      client.staticIp = undefined;
      client.forceVlan = false;
      client.vlanId = undefined;
      return;
    }

    client.useDhcp = client.useDhcp === true;
    if (client.useDhcp) {
      client.staticIp = undefined;
    }

    client.forceVlan = client.forceVlan === true;
    if (!client.forceVlan) {
      client.vlanId = undefined;
    }
  }

  private buildClientRuntimeUpdate(client: VpnClientDoc): Partial<plugins.smartvpn.IClientEntry> {
    return {
      description: client.description,
      security: this.buildClientSecurity(client),
      useHostIp: client.useHostIp,
      useDhcp: client.useDhcp,
      staticIp: client.staticIp,
      forceVlan: client.forceVlan,
      vlanId: client.vlanId,
    };
  }

  private async ensureForwardingModeForHostIpClient(useHostIp: boolean): Promise<void> {
    if (!useHostIp || !this.vpnServer) return;
    if (this.getResolvedForwardingMode() !== 'socket') return;

    logger.log('info', 'VPN: Restarting server in hybrid mode to support a host-IP client');
    this.forwardingModeOverride = 'hybrid';
    await this.stop();
    await this.start();
  }

  private async persistClient(client: VpnClientDoc): Promise<void> {
    await client.save();
  }
}