p2ns/includes/networking/dns_pool.js

const dgram = require('dgram');
const dnsPacket = require('dns-packet');
const { logDebug, logError, logWarn, logInfo } = require('../infrastructure/logger');

// DNS resolver connection pool
class DNSPool {
  constructor(maxConnections = 5) {
    this.maxConnections = maxConnections;
    this.resolvers = [];
    this.activeQueries = new Map();
    this.queryId = 0;
  }

  /**
   * Get or create a DNS resolver socket
   * @returns {dgram.Socket} - DNS resolver socket
   */
  getResolver() {
    if (this.resolvers.length < this.maxConnections) {
      const resolver = dgram.createSocket('udp4');
      resolver.on('error', (err) => {
        logError('DNSPool', `Resolver error: ${err.message}`);
      });
      this.resolvers.push(resolver);
      return resolver;
    }
    // Round-robin through existing resolvers
    return this.resolvers[this.queryId % this.resolvers.length];
  }

  /**
   * Get list of DNS servers from environment variable
   * @returns {string[]} - Array of DNS server IP addresses
   */
  getDnsServers() {
    const publicDnsServerEnv = process.env.PUBLIC_DNS_SERVER || '1.1.1.1';
    return publicDnsServerEnv.split(',').map(s => s.trim()).filter(s => s.length > 0);
  }

  /**
   * Query public DNS with failover support
   * @param {object} query - DNS query object
   * @returns {Promise<object|null>} - DNS response or null
   */
  async query(query) {
    const dnsServers = this.getDnsServers();
    
    // Try each DNS server in order (failover strategy)
    for (let serverIndex = 0; serverIndex < dnsServers.length; serverIndex++) {
      const publicDNSServer = dnsServers[serverIndex];
      const isLastServer = serverIndex === dnsServers.length - 1;
      
      const result = await new Promise((resolve) => {
        const resolver = this.getResolver();
        const queryId = ++this.queryId;
        let resolved = false;
        
        const cleanup = () => {
          if (!resolved) {
            resolved = true;
            const queryData = this.activeQueries.get(queryId);
            if (queryData) {
              // Remove listeners explicitly
              if (queryData.handler) {
                resolver.removeListener('message', queryData.handler);
              }
              if (queryData.errorHandler) {
                resolver.removeListener('error', queryData.errorHandler);
              }
              if (queryData.timeout) {
                clearTimeout(queryData.timeout);
              }
              this.activeQueries.delete(queryId);
            }
          }
        };
        
        // Set timeout
        const timeout = setTimeout(() => {
          logWarn('DNSPool', `DNS query timeout for ${query.questions[0].name} on server ${publicDNSServer}`);
          cleanup();
          resolve(null);
        }, 5000);
        
        const handler = (msg) => {
          if (resolved) return;
          clearTimeout(timeout);
          try {
            const response = dnsPacket.decode(msg);
            logDebug('DNSPool', `Public DNS response for ${query.questions[0].name} from ${publicDNSServer}: ${JSON.stringify(response.answers)}`);
            cleanup();
            resolve(response);
          } catch (err) {
            logError('DNSPool', `Error decoding DNS response from ${publicDNSServer}: ${err.message}`);
            cleanup();
            resolve(null);
          }
        };
        
        const errorHandler = (err) => {
          if (resolved) return;
          clearTimeout(timeout);
          logWarn('DNSPool', `Error receiving DNS response from ${publicDNSServer}: ${err.message}`);
          cleanup();
          resolve(null);
        };
        
        // Store query data with resolver reference for proper cleanup
        this.activeQueries.set(queryId, { handler, errorHandler, timeout, resolver });
        resolver.once('message', handler);
        resolver.once('error', errorHandler);
        
        const encodedQuery = dnsPacket.encode(query);
        resolver.send(encodedQuery, 53, publicDNSServer, (err) => {
          if (err) {
            clearTimeout(timeout);
            logWarn('DNSPool', `Error forwarding DNS query to ${publicDNSServer}: ${err.message}`);
            resolver.removeListener('message', handler);
            resolver.removeListener('error', errorHandler);
            cleanup();
            resolve(null);
          }
        });
      });
      
      // If we got a successful response, return it
      if (result) {
        return result;
      }
      
      // If this was the last server, return null
      if (isLastServer) {
        logWarn('DNSPool', `All DNS servers failed for ${query.questions[0].name}`);
        return null;
      }
      
      // Otherwise, try next server
      logDebug('DNSPool', `DNS server ${publicDNSServer} failed, trying next server...`);
    }
    
    return null;
  }

  /**
   * Close all resolvers
   */
  close() {
    // Clean up all active queries and their listeners
    for (const [queryId, query] of this.activeQueries.entries()) {
      if (query.timeout) {
        clearTimeout(query.timeout);
      }
      // Remove listeners from resolver
      if (query.resolver && query.handler) {
        try {
          query.resolver.removeListener('message', query.handler);
        } catch (err) {
          logDebug('DNSPool', `Error removing message listener: ${err.message}`);
        }
      }
      if (query.resolver && query.errorHandler) {
        try {
          query.resolver.removeListener('error', query.errorHandler);
        } catch (err) {
          logDebug('DNSPool', `Error removing error listener: ${err.message}`);
        }
      }
    }
    this.activeQueries.clear();
    
    // Close all resolvers and remove all listeners
    for (const resolver of this.resolvers) {
      try {
        // Remove all listeners to prevent leaks
        resolver.removeAllListeners();
        resolver.close();
      } catch (err) {
        logError('DNSPool', `Error closing resolver: ${err.message}`);
      }
    }
    this.resolvers = [];
  }
}

// Global DNS pool instance
const dnsPool = new DNSPool(parseInt(process.env.DNS_POOL_SIZE || '5', 10));

/**
 * Update DNS pool configuration at runtime
 * @param {number} maxConnections - New maximum number of connections
 */
function updateDnsPool(maxConnections) {
  const oldMax = dnsPool.maxConnections;
  dnsPool.maxConnections = maxConnections;
  
  // If reducing pool size, close excess resolvers
  if (maxConnections < oldMax && dnsPool.resolvers.length > maxConnections) {
    const excessResolvers = dnsPool.resolvers.splice(maxConnections);
    for (const resolver of excessResolvers) {
      try {
        resolver.removeAllListeners();
        resolver.close();
      } catch (err) {
        logError('DNSPool', `Error closing excess resolver: ${err.message}`);
      }
    }
    logInfo('DNSPool', `Reduced pool size from ${oldMax} to ${maxConnections}, closed ${excessResolvers.length} resolvers`);
  } else if (maxConnections > oldMax) {
    logInfo('DNSPool', `Increased pool size from ${oldMax} to ${maxConnections}, new resolvers will be created on demand`);
  }
  
  // PUBLIC_DNS_SERVER is already read from process.env at query time, so no action needed
}

module.exports = { dnsPool, DNSPool, updateDnsPool };