smartdns/ts_server/classes.dnsserver.ts

import * as plugins from './plugins.js';
import { RustDnsBridge, type IDnsQueryEvent, type IIpcDnsAnswer, type IIpcDnsQuestion, type IRustDnsConfig } from './classes.rustdnsbridge.js';

export interface IDnsServerOptions {
  httpsKey: string;
  httpsCert: string;
  httpsPort: number;
  udpPort: number;
  dnssecZone: string;
  udpBindInterface?: string;
  httpsBindInterface?: string;
  // New options for independent manual socket control
  manualUdpMode?: boolean;
  manualHttpsMode?: boolean;
  // Primary nameserver for SOA records (defaults to ns1.{dnssecZone})
  primaryNameserver?: string;
  // Local handling for RFC 6761 localhost (default: true)
  enableLocalhostHandling?: boolean;
}

export interface DnsAnswer {
  name: string;
  type: string;
  class: string | number;
  ttl: number;
  data: any;
}

export interface IDnsQuestion {
  name: string;
  type: string;
  class?: string;
}

export interface IDnsHandler {
  domainPattern: string;
  recordTypes: string[];
  handler: (question: IDnsQuestion) => DnsAnswer | null;
}

// Let's Encrypt related interfaces
interface LetsEncryptOptions {
  email?: string;
  staging?: boolean;
  certDir?: string;
}

export interface IDnsQueryCompletedEvent {
  /** The original questions from the query */
  questions: IIpcDnsQuestion[];
  /** Whether any handler answered the query */
  answered: boolean;
  /** How long handler resolution took (ms) */
  responseTimeMs: number;
  /** Timestamp of the query */
  timestamp: number;
}

export class DnsServer extends plugins.events.EventEmitter {
  private bridge: RustDnsBridge;
  private handlers: IDnsHandler[] = [];

  // Track initialization state
  private bridgeSpawned: boolean = false;

  // Legacy server references (kept for backward-compatible test access)
  private httpsServer: any = null;
  private udpServer: any = null;

  constructor(private options: IDnsServerOptions) {
    super();
    this.bridge = new RustDnsBridge();

    // Wire up the dnsQuery event to run TypeScript handlers
    this.bridge.on('dnsQuery', async (event: IDnsQueryEvent) => {
      try {
        const startTime = Date.now();
        const answers = this.resolveQuery(event);
        const responseTimeMs = Date.now() - startTime;
        this.emit('query', {
          questions: event.questions,
          answered: answers.answered,
          responseTimeMs,
          timestamp: startTime,
        } satisfies IDnsQueryCompletedEvent);
        await this.bridge.sendQueryResult(
          event.correlationId,
          answers.answers,
          answers.answered
        );
      } catch (err) {
        console.error('Error handling DNS query:', err);
        try {
          await this.bridge.sendQueryResult(event.correlationId, [], false);
        } catch (sendErr) {
          console.error('Error sending empty query result:', sendErr);
        }
      }
    });
  }

  /**
   * Register a DNS handler for a domain pattern and record types.
   */
  public registerHandler(
    domainPattern: string,
    recordTypes: string[],
    handler: (question: IDnsQuestion) => DnsAnswer | null
  ): void {
    this.handlers.push({ domainPattern, recordTypes, handler });
  }

  /**
   * Unregister a specific handler.
   */
  public unregisterHandler(domainPattern: string, recordTypes: string[]): boolean {
    const initialLength = this.handlers.length;
    this.handlers = this.handlers.filter(handler =>
      !(handler.domainPattern === domainPattern &&
        recordTypes.every(type => handler.recordTypes.includes(type)))
    );
    return this.handlers.length < initialLength;
  }

  /**
   * Start the DNS server.
   */
  public async start(): Promise<void> {
    // Validate interface addresses if provided
    const udpInterface = this.options.udpBindInterface || '0.0.0.0';
    const httpsInterface = this.options.httpsBindInterface || '0.0.0.0';

    if (this.options.udpBindInterface && !this.isValidIpAddress(this.options.udpBindInterface)) {
      throw new Error(`Invalid UDP bind interface: ${this.options.udpBindInterface}`);
    }

    if (this.options.httpsBindInterface && !this.isValidIpAddress(this.options.httpsBindInterface)) {
      throw new Error(`Invalid HTTPS bind interface: ${this.options.httpsBindInterface}`);
    }

    // Spawn the Rust binary
    if (!this.bridgeSpawned) {
      const spawned = await this.bridge.spawn();
      if (!spawned) {
        throw new Error('Failed to spawn rustdns binary');
      }
      this.bridgeSpawned = true;
    }

    // Build config for Rust
    const config: IRustDnsConfig = {
      udpPort: this.options.udpPort,
      httpsPort: this.options.httpsPort,
      udpBindInterface: udpInterface,
      httpsBindInterface: httpsInterface,
      httpsKey: this.options.httpsKey || '',
      httpsCert: this.options.httpsCert || '',
      dnssecZone: this.options.dnssecZone,
      dnssecAlgorithm: 'ECDSA',
      primaryNameserver: this.options.primaryNameserver || `ns1.${this.options.dnssecZone}`,
      enableLocalhostHandling: this.options.enableLocalhostHandling !== false,
      manualUdpMode: this.options.manualUdpMode || false,
      manualHttpsMode: this.options.manualHttpsMode || false,
    };

    await this.bridge.startServer(config);

    // Set legacy markers for backward-compatible test checks
    this.httpsServer = { _rustBridgeManaged: true };
    this.udpServer = { _rustBridgeManaged: true };

    const udpManual = this.options.manualUdpMode || false;
    const httpsManual = this.options.manualHttpsMode || false;

    if (udpManual && httpsManual) {
      console.log('DNS server started in full manual mode - ready to accept connections');
    } else if (udpManual && !httpsManual) {
      console.log('DNS server started with manual UDP mode and automatic HTTPS binding');
    } else if (!udpManual && httpsManual) {
      console.log('DNS server started with automatic UDP binding and manual HTTPS mode');
    } else {
      console.log(`DNS server started (UDP: ${udpInterface}:${this.options.udpPort}, HTTPS: ${httpsInterface}:${this.options.httpsPort})`);
    }
  }

  /**
   * Stop the DNS server.
   */
  public async stop(): Promise<void> {
    if (this.bridgeSpawned) {
      try {
        await this.bridge.stopServer();
      } catch (err) {
        // Ignore errors during stop (process may have already exited)
      }
      this.bridge.kill();
      this.bridgeSpawned = false;
    }

    this.httpsServer = null;
    this.udpServer = null;
  }

  /**
   * Initialize servers (no-op with Rust bridge, kept for API compatibility).
   */
  public initializeServers(): void {
    // No-op — Rust bridge handles server initialization via start()
  }

  /**
   * Initialize UDP server (no-op with Rust bridge).
   */
  public initializeUdpServer(): void {
    // No-op
  }

  /**
   * Initialize HTTPS server (no-op with Rust bridge).
   */
  public initializeHttpsServer(): void {
    // No-op
  }

  /**
   * Handle a raw TCP socket for HTTPS/DoH.
   * In Rust mode, this is not directly supported — use processRawDnsPacket instead.
   */
  public handleHttpsSocket(socket: plugins.net.Socket): void {
    console.warn('handleHttpsSocket: direct socket handling not available with Rust bridge. Use processRawDnsPacket instead.');
  }

  /**
   * Handle a UDP message manually.
   */
  public handleUdpMessage(
    msg: Buffer,
    rinfo: plugins.dgram.RemoteInfo,
    responseCallback?: (response: Buffer, rinfo: plugins.dgram.RemoteInfo) => void
  ): void {
    // Async processing via Rust bridge
    this.processRawDnsPacketAsync(msg)
      .then((responseData) => {
        if (responseCallback) {
          responseCallback(responseData, rinfo);
        }
      })
      .catch((err) => {
        console.error('Error processing UDP DNS request:', err);
      });
  }

  /**
   * Process a raw DNS packet asynchronously via Rust bridge.
   */
  public async processRawDnsPacketAsync(packet: Buffer): Promise<Buffer> {
    if (!this.bridgeSpawned) {
      throw new Error('DNS server not started — call start() first');
    }
    return this.bridge.processPacket(packet);
  }

  /**
   * Retrieve SSL certificate for specified domains using Let's Encrypt
   */
  public async retrieveSslCertificate(
    domainNames: string[],
    options: LetsEncryptOptions = {}
  ): Promise<{ cert: string; key: string; success: boolean }> {
    const opts = {
      email: options.email || 'admin@example.com',
      staging: options.staging !== undefined ? options.staging : false,
      certDir: options.certDir || './certs'
    };

    if (!plugins.fs.existsSync(opts.certDir)) {
      plugins.fs.mkdirSync(opts.certDir, { recursive: true });
    }

    const authorizedDomains = this.filterAuthorizedDomains(domainNames);

    if (authorizedDomains.length === 0) {
      console.error('None of the provided domains are authorized for this DNS server');
      return { cert: '', key: '', success: false };
    }

    console.log(`Retrieving SSL certificate for domains: ${authorizedDomains.join(', ')}`);

    try {
      // @ts-ignore - acmeClientOverride is added for testing purposes
      const acmeClient = this.acmeClientOverride || await import('acme-client');

      const accountKeyPath = plugins.path.join(opts.certDir, 'account.key');
      let accountKey: Buffer;

      if (plugins.fs.existsSync(accountKeyPath)) {
        accountKey = plugins.fs.readFileSync(accountKeyPath);
      } else {
        const { privateKey } = plugins.crypto.generateKeyPairSync('rsa', {
          modulusLength: 2048,
          publicKeyEncoding: { type: 'spki', format: 'pem' },
          privateKeyEncoding: { type: 'pkcs8', format: 'pem' }
        });

        accountKey = Buffer.from(privateKey);
        plugins.fs.writeFileSync(accountKeyPath, accountKey);
      }

      const client = new acmeClient.Client({
        directoryUrl: opts.staging
          ? acmeClient.directory.letsencrypt.staging
          : acmeClient.directory.letsencrypt.production,
        accountKey: accountKey
      });

      await client.createAccount({
        termsOfServiceAgreed: true,
        contact: [`mailto:${opts.email}`]
      });

      const order = await client.createOrder({
        identifiers: authorizedDomains.map(domain => ({
          type: 'dns',
          value: domain
        }))
      });

      const authorizations = await client.getAuthorizations(order);
      const challengeHandlers: { domain: string; pattern: string }[] = [];

      for (const auth of authorizations) {
        const domain = auth.identifier.value;
        const challenge = auth.challenges.find((c: any) => c.type === 'dns-01');
        if (!challenge) {
          throw new Error(`No DNS-01 challenge found for ${domain}`);
        }

        const keyAuthorization = await client.getChallengeKeyAuthorization(challenge);
        const recordValue = this.getDnsRecordValueForChallenge(keyAuthorization);
        const challengeDomain = `_acme-challenge.${domain}`;

        console.log(`Setting up TXT record for ${challengeDomain}: ${recordValue}`);

        this.registerHandler(
          challengeDomain,
          ['TXT'],
          (question: IDnsQuestion): DnsAnswer | null => {
            if (question.name === challengeDomain && question.type === 'TXT') {
              return {
                name: question.name,
                type: 'TXT',
                class: 'IN',
                ttl: 300,
                data: [recordValue]
              };
            }
            return null;
          }
        );

        challengeHandlers.push({ domain, pattern: challengeDomain });
        await new Promise(resolve => setTimeout(resolve, 2000));
        await client.completeChallenge(challenge);
        await client.waitForValidStatus(challenge);
        console.log(`Challenge for ${domain} validated successfully!`);
      }

      const domainKeyPath = plugins.path.join(opts.certDir, `${authorizedDomains[0]}.key`);
      const { privateKey } = plugins.crypto.generateKeyPairSync('rsa', {
        modulusLength: 2048,
        publicKeyEncoding: { type: 'spki', format: 'pem' },
        privateKeyEncoding: { type: 'pkcs8', format: 'pem' }
      });

      plugins.fs.writeFileSync(domainKeyPath, privateKey);

      interface CSRResult {
        csr: Buffer;
      }

      const csrResult = await acmeClient.forge.createCsr({
        commonName: authorizedDomains[0],
        altNames: authorizedDomains
      }) as unknown as CSRResult;

      await client.finalizeOrder(order, csrResult.csr);
      const certificate = await client.getCertificate(order);

      const certPath = plugins.path.join(opts.certDir, `${authorizedDomains[0]}.cert`);
      plugins.fs.writeFileSync(certPath, certificate);

      this.options.httpsCert = certificate;
      this.options.httpsKey = privateKey;

      // Update certs on Rust bridge if running
      if (this.bridgeSpawned) {
        try {
          await this.bridge.updateCerts(privateKey, certificate);
        } catch (err) {
          console.error('Error updating certs on Rust bridge:', err);
        }
      }

      for (const handler of challengeHandlers) {
        this.unregisterHandler(handler.pattern, ['TXT']);
        console.log(`Cleaned up challenge handler for ${handler.domain}`);
      }

      return {
        cert: certificate,
        key: privateKey,
        success: true
      };
    } catch (error) {
      console.error('Error retrieving SSL certificate:', error);
      return { cert: '', key: '', success: false };
    }
  }

  /**
   * Filter domains to include only those the server is authoritative for.
   */
  public filterAuthorizedDomains(domainNames: string[]): string[] {
    const authorizedDomains: string[] = [];

    for (const domain of domainNames) {
      if (domain.startsWith('*.')) {
        const baseDomain = domain.substring(2);
        if (this.isAuthorizedForDomain(baseDomain)) {
          authorizedDomains.push(domain);
        }
      } else if (this.isAuthorizedForDomain(domain)) {
        authorizedDomains.push(domain);
      }
    }

    return authorizedDomains;
  }

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

  /**
   * Resolve a DNS query event from Rust using TypeScript handlers.
   */
  private resolveQuery(event: IDnsQueryEvent): { answers: IIpcDnsAnswer[]; answered: boolean } {
    const answers: IIpcDnsAnswer[] = [];
    let answered = false;

    for (const q of event.questions) {
      const question: IDnsQuestion = {
        name: q.name,
        type: q.type,
        class: q.class,
      };

      for (const handlerEntry of this.handlers) {
        if (
          plugins.minimatch.minimatch(q.name, handlerEntry.domainPattern) &&
          handlerEntry.recordTypes.includes(q.type)
        ) {
          const answer = handlerEntry.handler(question);
          if (answer) {
            answers.push({
              name: answer.name,
              type: answer.type,
              class: typeof answer.class === 'number' ? 'IN' : (answer.class || 'IN'),
              ttl: answer.ttl || 300,
              data: answer.data,
            });
            answered = true;
          }
        }
      }
    }

    return { answers, answered };
  }

  private getDnsRecordValueForChallenge(keyAuthorization: string): string {
    const digest = plugins.crypto
      .createHash('sha256')
      .update(keyAuthorization)
      .digest('base64')
      .replace(/\+/g, '-')
      .replace(/\//g, '_')
      .replace(/=/g, '');

    return digest;
  }

  private isValidIpAddress(ip: string): boolean {
    const ipv4Pattern = /^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$/;
    const ipv6Pattern = /^(::1|::)$|^([0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}$/;
    return ipv4Pattern.test(ip) || ipv6Pattern.test(ip);
  }

  private isAuthorizedForDomain(domain: string): boolean {
    for (const handler of this.handlers) {
      if (plugins.minimatch.minimatch(domain, handler.domainPattern)) {
        return true;
      }
    }

    if (domain === this.options.dnssecZone || domain.endsWith(`.${this.options.dnssecZone}`)) {
      return true;
    }

    return false;
  }
}