smartdns/test/test.server.ts

import * as plugins from '../ts_server/plugins.js';
import { expect, tap } from '@push.rocks/tapbundle';
import { tapNodeTools } from '@push.rocks/tapbundle/node';
import { execSync } from 'child_process';

import * as dnsPacket from 'dns-packet';
import * as https from 'https';
import * as dgram from 'dgram';
import * as fs from 'fs';
import * as path from 'path';
import * as os from 'os';

import * as smartdns from '../ts_server/index.js';

// Generate a real self-signed certificate using OpenSSL
function generateSelfSignedCert() {
  const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'cert-'));
  const keyPath = path.join(tmpDir, 'key.pem');
  const certPath = path.join(tmpDir, 'cert.pem');
  
  try {
    // Generate private key
    execSync(`openssl genrsa -out "${keyPath}" 2048`);
    
    // Generate self-signed certificate
    execSync(
      `openssl req -new -x509 -key "${keyPath}" -out "${certPath}" -days 365 -subj "/C=US/ST=State/L=City/O=Organization/CN=test.example.com"`
    );
    
    // Read the files
    const privateKey = fs.readFileSync(keyPath, 'utf8');
    const cert = fs.readFileSync(certPath, 'utf8');
    
    return { key: privateKey, cert };
  } catch (error) {
    console.error('Error generating certificate:', error);
    throw error;
  } finally {
    // Clean up temporary files
    try {
      if (fs.existsSync(keyPath)) fs.unlinkSync(keyPath);
      if (fs.existsSync(certPath)) fs.unlinkSync(certPath);
      if (fs.existsSync(tmpDir)) fs.rmdirSync(tmpDir);
    } catch (err) {
      console.error('Error cleaning up temporary files:', err);
    }
  }
}

// Cache the generated certificate for performance
let cachedCert = null;

// Helper function to get certificate
function getTestCertificate() {
  if (!cachedCert) {
    cachedCert = generateSelfSignedCert();
  }
  return cachedCert;
}

// Mock for acme-client directly imported as a module
const acmeClientMock = {
  Client: class {
    constructor() {}
    
    createAccount() {
      return Promise.resolve({});
    }
    
    createOrder() {
      return Promise.resolve({
        authorizations: ['auth1', 'auth2']
      });
    }
    
    getAuthorizations() {
      return Promise.resolve([
        {
          identifier: { value: 'test.bleu.de' },
          challenges: [
            { type: 'dns-01', url: 'https://example.com/challenge' }
          ]
        }
      ]);
    }
    
    getChallengeKeyAuthorization() {
      return Promise.resolve('test_key_authorization');
    }
    
    completeChallenge() {
      return Promise.resolve({});
    }
    
    waitForValidStatus() {
      return Promise.resolve({});
    }
    
    finalizeOrder() {
      return Promise.resolve({});
    }
    
    getCertificate() {
      // Use a real certificate
      const { cert } = getTestCertificate();
      return Promise.resolve(cert);
    }
  },
  
  forge: {
    createCsr({commonName, altNames}) {
      return Promise.resolve({
        csr: Buffer.from('mock-csr-data')
      });
    }
  },
  
  directory: {
    letsencrypt: {
      staging: 'https://acme-staging-v02.api.letsencrypt.org/directory',
      production: 'https://acme-v02.api.letsencrypt.org/directory'
    }
  }
};

// Override generateKeyPairSync to use our test key for certificate generation in tests
const originalGenerateKeyPairSync = plugins.crypto.generateKeyPairSync;
plugins.crypto.generateKeyPairSync = function(type, options) {
  if (type === 'rsa' && 
      options?.modulusLength === 2048 && 
      options?.privateKeyEncoding?.type === 'pkcs8') {
    
    // Get the test certificate key if we're in the retrieveSslCertificate method
    try {
      const stack = new Error().stack || '';
      if (stack.includes('retrieveSslCertificate')) {
        const { key } = getTestCertificate();
        return { privateKey: key, publicKey: 'TEST_PUBLIC_KEY' };
      }
    } catch (e) {
      // Fall back to original function if error occurs
    }
  }
  
  // Use the original function for other cases
  return originalGenerateKeyPairSync.apply(this, arguments);
};

let dnsServer: smartdns.DnsServer;
const testCertDir = path.join(process.cwd(), 'test-certs');

// Helper to clean up test certificate directory
function cleanCertDir() {
  if (fs.existsSync(testCertDir)) {
    const files = fs.readdirSync(testCertDir);
    for (const file of files) {
      fs.unlinkSync(path.join(testCertDir, file));
    }
    fs.rmdirSync(testCertDir);
  }
}

// Port management for tests
let nextHttpsPort = 8080;
let nextUdpPort = 8081;

function getUniqueHttpsPort() {
  return nextHttpsPort++;
}

function getUniqueUdpPort() {
  return nextUdpPort++;
}

// Cleanup function for servers - more robust implementation
async function stopServer(server: smartdns.DnsServer | null | undefined) {
  if (!server) {
    return; // Nothing to do if server doesn't exist
  }
  
  try {
    // Access private properties for checking before stopping
    // @ts-ignore - accessing private properties for testing
    const hasHttpsServer = server.httpsServer !== undefined && server.httpsServer !== null;
    // @ts-ignore - accessing private properties for testing
    const hasUdpServer = server.udpServer !== undefined && server.udpServer !== null;
    
    // Only try to stop if there's something to stop
    if (hasHttpsServer || hasUdpServer) {
      await server.stop();
    }
  } catch (e) {
    console.log('Handled error when stopping server:', e);
    // Ignore errors during cleanup
  }
}

// Setup and teardown
tap.test('setup', async () => {
  cleanCertDir();
  // Reset dnsServer to null at the start
  dnsServer = null;
  // Reset certificate cache
  cachedCert = null;
});

tap.test('teardown', async () => {
  // Stop the server if it exists
  await stopServer(dnsServer);
  dnsServer = null;
  
  cleanCertDir();
  // Reset certificate cache
  cachedCert = null;
});

tap.test('should create an instance of DnsServer', async () => {
  // Use valid options
  const httpsData = await tapNodeTools.createHttpsCert();
  dnsServer = new smartdns.DnsServer({
    httpsKey: httpsData.key,
    httpsCert: httpsData.cert,
    httpsPort: 8080,
    udpPort: 8081,
    dnssecZone: 'example.com',
  });
  expect(dnsServer).toBeInstanceOf(smartdns.DnsServer);
});

tap.test('should start the server', async () => {
  // Clean up any existing server
  await stopServer(dnsServer);
  
  const httpsData = await tapNodeTools.createHttpsCert();
  dnsServer = new smartdns.DnsServer({
    httpsKey: httpsData.key,
    httpsCert: httpsData.cert,
    httpsPort: getUniqueHttpsPort(),
    udpPort: getUniqueUdpPort(),
    dnssecZone: 'example.com',
  });
  
  await dnsServer.start();
  // @ts-ignore - accessing private property for testing
  expect(dnsServer.httpsServer).toBeDefined();
  
  // Stop the server at the end of this test
  await stopServer(dnsServer);
  dnsServer = null;
});

tap.test('lets add a handler', async () => {
  const httpsData = await tapNodeTools.createHttpsCert();
  dnsServer = new smartdns.DnsServer({
    httpsKey: httpsData.key,
    httpsCert: httpsData.cert,
    httpsPort: 8080,
    udpPort: 8081,
    dnssecZone: 'example.com',
  });

  dnsServer.registerHandler('*.bleu.de', ['A'], (question) => {
    return {
      name: question.name,
      type: 'A',
      class: 'IN',
      ttl: 300,
      data: '127.0.0.1',
    };
  });

  // @ts-ignore - accessing private method for testing
  const response = dnsServer.processDnsRequest({
    type: 'query',
    id: 1,
    flags: 0,
    questions: [
      {
        name: 'dnsly_a.bleu.de',
        type: 'A',
        class: 'IN',
      },
    ],
    answers: [],
  });
  expect(response.answers[0]).toEqual({
    name: 'dnsly_a.bleu.de',
    type: 'A',
    class: 'IN',
    ttl: 300,
    data: '127.0.0.1',
  });
});

tap.test('should unregister a handler', async () => {
  const httpsData = await tapNodeTools.createHttpsCert();
  dnsServer = new smartdns.DnsServer({
    httpsKey: httpsData.key,
    httpsCert: httpsData.cert,
    httpsPort: 8080,
    udpPort: 8081,
    dnssecZone: 'example.com',
  });

  // Register handlers
  dnsServer.registerHandler('*.bleu.de', ['A'], (question) => {
    return {
      name: question.name,
      type: 'A',
      class: 'IN',
      ttl: 300,
      data: '127.0.0.1',
    };
  });
  
  dnsServer.registerHandler('test.com', ['TXT'], (question) => {
    return {
      name: question.name,
      type: 'TXT',
      class: 'IN',
      ttl: 300,
      data: ['test'],
    };
  });

  // Test unregistering
  const result = dnsServer.unregisterHandler('*.bleu.de', ['A']);
  expect(result).toEqual(true);
  
  // Verify handler is removed
  // @ts-ignore - accessing private method for testing
  const response = dnsServer.processDnsRequest({
    type: 'query',
    id: 1,
    flags: 0,
    questions: [
      {
        name: 'dnsly_a.bleu.de',
        type: 'A',
        class: 'IN',
      },
    ],
    answers: [],
  });
  
  // Should get SOA record instead of A record
  expect(response.answers[0].type).toEqual('SOA');
});

tap.test('lets query over https', async () => {
  // Clean up any existing server
  await stopServer(dnsServer);
  
  const httpsPort = getUniqueHttpsPort();
  const httpsData = await tapNodeTools.createHttpsCert();
  
  dnsServer = new smartdns.DnsServer({
    httpsKey: httpsData.key,
    httpsCert: httpsData.cert,
    httpsPort: httpsPort,
    udpPort: getUniqueUdpPort(),
    dnssecZone: 'example.com',
  });
  
  await dnsServer.start();
  
  dnsServer.registerHandler('*.bleu.de', ['A'], (question) => {
    return {
      name: question.name,
      type: 'A',
      class: 'IN',
      ttl: 300,
      data: '127.0.0.1',
    };
  });

  // Skip SSL verification for self-signed cert in tests
  process.env.NODE_TLS_REJECT_UNAUTHORIZED = '0';
  
  const query = dnsPacket.encode({
    type: 'query',
    id: 2,
    flags: dnsPacket.RECURSION_DESIRED,
    questions: [
      {
        name: 'dnsly_a.bleu.de',
        type: 'A',
        class: 'IN',
      },
    ],
  });

  const response = await fetch(`https://localhost:${httpsPort}/dns-query`, {
    method: 'POST',
    body: query,
    headers: {
      'Content-Type': 'application/dns-message',
    }
  });

  expect(response.status).toEqual(200);

  const responseData = await response.arrayBuffer();
  const dnsResponse = dnsPacket.decode(Buffer.from(responseData));

  console.log(dnsResponse.answers[0]);

  expect(dnsResponse.answers[0]).toEqual({
    name: 'dnsly_a.bleu.de',
    type: 'A',
    class: 'IN',
    ttl: 300,
    flush: false,
    data: '127.0.0.1',
  });
  
  // Reset TLS verification
  process.env.NODE_TLS_REJECT_UNAUTHORIZED = '1';
  
  // Clean up server
  await stopServer(dnsServer);
  dnsServer = null;
});

tap.test('lets query over udp', async () => {
  // Clean up any existing server
  await stopServer(dnsServer);
  
  const udpPort = getUniqueUdpPort();
  const httpsData = await tapNodeTools.createHttpsCert();
  
  dnsServer = new smartdns.DnsServer({
    httpsKey: httpsData.key,
    httpsCert: httpsData.cert,
    httpsPort: getUniqueHttpsPort(),
    udpPort: udpPort,
    dnssecZone: 'example.com',
  });
  
  await dnsServer.start();
  
  dnsServer.registerHandler('*.bleu.de', ['A'], (question) => {
    return {
      name: question.name,
      type: 'A',
      class: 'IN',
      ttl: 300,
      data: '127.0.0.1',
    };
  });

  const client = dgram.createSocket('udp4');

  const query = dnsPacket.encode({
    type: 'query',
    id: 3,
    flags: dnsPacket.RECURSION_DESIRED,
    questions: [
      {
        name: 'dnsly_a.bleu.de',
        type: 'A',
        class: 'IN',
      },
    ],
  });

  const responsePromise = new Promise<dnsPacket.Packet>((resolve, reject) => {
    client.on('message', (msg) => {
      const dnsResponse = dnsPacket.decode(msg);
      resolve(dnsResponse);
      client.close();
    });

    client.on('error', (err) => {
      reject(err);
      client.close();
    });

    client.send(query, udpPort, 'localhost', (err) => {
      if (err) {
        reject(err);
        client.close();
      }
    });
  });

  const dnsResponse = await responsePromise;

  console.log(dnsResponse.answers[0]);

  expect(dnsResponse.answers[0]).toEqual({
    name: 'dnsly_a.bleu.de',
    type: 'A',
    class: 'IN',
    ttl: 300,
    flush: false,
    data: '127.0.0.1',
  });
  
  // Clean up server
  await stopServer(dnsServer);
  dnsServer = null;
});

tap.test('should filter authorized domains correctly', async () => {
  const httpsData = await tapNodeTools.createHttpsCert();
  dnsServer = new smartdns.DnsServer({
    httpsKey: httpsData.key,
    httpsCert: httpsData.cert,
    httpsPort: 8080,
    udpPort: 8081,
    dnssecZone: 'example.com',
  });
  
  // Register handlers for specific domains
  dnsServer.registerHandler('*.bleu.de', ['A'], () => null);
  dnsServer.registerHandler('test.com', ['A'], () => null);
  
  // Test filtering authorized domains
  const authorizedDomains = dnsServer.filterAuthorizedDomains([
    'test.com',              // Should be authorized
    'sub.test.com',          // Should not be authorized
    '*.bleu.de',             // Pattern itself isn't a domain
    'something.bleu.de',     // Should be authorized via wildcard pattern
    'example.com',           // Should be authorized (dnssecZone)
    'sub.example.com',       // Should be authorized (within dnssecZone)
    'othersite.org'          // Should not be authorized
  ]);
  
  // Using toContain with expect from tapbundle
  expect(authorizedDomains.includes('test.com')).toEqual(true);
  expect(authorizedDomains.includes('something.bleu.de')).toEqual(true);
  expect(authorizedDomains.includes('example.com')).toEqual(true);
  expect(authorizedDomains.includes('sub.example.com')).toEqual(true);
  
  expect(authorizedDomains.includes('sub.test.com')).toEqual(false);
  expect(authorizedDomains.includes('*.bleu.de')).toEqual(false);
  expect(authorizedDomains.includes('othersite.org')).toEqual(false);
});

tap.test('should retrieve SSL certificate successfully', async () => {
  // Clean up any existing server
  await stopServer(dnsServer);
  
  // Create a temporary directory for the certificate test
  const tempCertDir = path.join(process.cwd(), 'temp-certs');
  if (!fs.existsSync(tempCertDir)) {
    fs.mkdirSync(tempCertDir, { recursive: true });
  }
  
  // Create a server with unique ports
  const httpsData = await tapNodeTools.createHttpsCert();
  dnsServer = new smartdns.DnsServer({
    httpsKey: httpsData.key,
    httpsCert: httpsData.cert,
    httpsPort: getUniqueHttpsPort(),
    udpPort: getUniqueUdpPort(),
    dnssecZone: 'example.com',
  });
  
  // Register handlers for test domains
  dnsServer.registerHandler('*.bleu.de', ['A'], () => null);
  dnsServer.registerHandler('test.bleu.de', ['A'], () => null);
  
  await dnsServer.start();
  
  // Inject our mock for acme-client
  (dnsServer as any).acmeClientOverride = acmeClientMock;
  
  try {
    // Request certificate for domains
    const result = await dnsServer.retrieveSslCertificate(
      ['test.bleu.de', '*.bleu.de', 'unknown.org'],
      {
        email: 'test@example.com',
        staging: true,
        certDir: tempCertDir
      }
    );
    
    console.log('Certificate retrieval result:', {
      success: result.success,
      certLength: result.cert.length,
      keyLength: result.key.length,
    });
    
    expect(result.success).toEqual(true);
    expect(result.cert.includes('BEGIN CERTIFICATE')).toEqual(true);
    expect(typeof result.key === 'string').toEqual(true);
    
    // Check that certificate directory was created
    expect(fs.existsSync(tempCertDir)).toEqual(true);
    
    // Verify TXT record handler was registered and then removed
    // @ts-ignore - accessing private property for testing
    const txtHandlerCount = dnsServer.handlers.filter(h => 
      h.domainPattern.includes('_acme-challenge') && 
      h.recordTypes.includes('TXT')
    ).length;
    
    expect(txtHandlerCount).toEqual(0); // Should be removed after validation
  } catch (err) {
    console.error('Test error:', err);
    throw err;
  } finally {
    // Clean up server and temporary cert directory
    await stopServer(dnsServer);
    dnsServer = null;
    
    if (fs.existsSync(tempCertDir)) {
      const files = fs.readdirSync(tempCertDir);
      for (const file of files) {
        fs.unlinkSync(path.join(tempCertDir, file));
      }
      fs.rmdirSync(tempCertDir);
    }
  }
});

tap.test('should run for a while', async (toolsArg) => {
  await toolsArg.delayFor(1000);
});

tap.test('should stop the server', async () => {
  // Clean up any existing server
  await stopServer(dnsServer);
  
  const httpsData = await tapNodeTools.createHttpsCert();
  dnsServer = new smartdns.DnsServer({
    httpsKey: httpsData.key,
    httpsCert: httpsData.cert,
    httpsPort: getUniqueHttpsPort(),
    udpPort: getUniqueUdpPort(),
    dnssecZone: 'example.com',
  });
  
  await dnsServer.start();
  await dnsServer.stop();
  
  // @ts-ignore - accessing private property for testing
  expect(dnsServer.httpsServer).toEqual(null);
  
  // Clear the reference
  dnsServer = null;
});

await tap.start();