test: add DNS restriction enforcement tests

src/cli.test.ts

@@ -48,14 +48,14 @@
 
     afterEach(() => {
       // Clean up the test directory
       if (fs.existsSync(testDir)) {
         fs.rmSync(testDir, { recursive: true, force: true });
       }
     });
 
     it('should parse domains from file with one domain per line', () => {
       const filePath = path.join(testDir, 'domains.txt');
       fs.writeFileSync(filePath, 'github.com\napi.github.com\nnpmjs.org');
 
       const result = parseDomainsFile(filePath);
 
@@ -64,7 +64,7 @@
 
     it('should parse comma-separated domains from file', () => {
       const filePath = path.join(testDir, 'domains.txt');
       fs.writeFileSync(filePath, 'github.com, api.github.com, npmjs.org');
 
       const result = parseDomainsFile(filePath);
 
@@ -73,7 +73,7 @@
 
     it('should handle mixed formats (lines and commas)', () => {
       const filePath = path.join(testDir, 'domains.txt');
       fs.writeFileSync(filePath, 'github.com\napi.github.com, npmjs.org\nexample.com');
 
       const result = parseDomainsFile(filePath);
 
@@ -82,7 +82,7 @@
 
     it('should skip empty lines', () => {
       const filePath = path.join(testDir, 'domains.txt');
       fs.writeFileSync(filePath, 'github.com\n\n\napi.github.com\n\nnpmjs.org');
 
       const result = parseDomainsFile(filePath);
 
@@ -91,7 +91,7 @@
 
     it('should skip lines with only whitespace', () => {
       const filePath = path.join(testDir, 'domains.txt');
       fs.writeFileSync(filePath, 'github.com\n   \n\t\napi.github.com');
 
       const result = parseDomainsFile(filePath);
 
@@ -100,7 +100,7 @@
 
     it('should skip comments starting with #', () => {
       const filePath = path.join(testDir, 'domains.txt');
       fs.writeFileSync(filePath, '# This is a comment\ngithub.com\n# Another comment\napi.github.com');
 
       const result = parseDomainsFile(filePath);
 
@@ -109,7 +109,7 @@
 
     it('should handle inline comments (after domain)', () => {
       const filePath = path.join(testDir, 'domains.txt');
       fs.writeFileSync(filePath, 'github.com # GitHub main domain\napi.github.com # API endpoint');
 
       const result = parseDomainsFile(filePath);
 
@@ -118,7 +118,7 @@
 
     it('should handle domains with inline comments in comma-separated format', () => {
       const filePath = path.join(testDir, 'domains.txt');
       fs.writeFileSync(filePath, 'github.com, api.github.com # GitHub domains\nnpmjs.org');
 
       const result = parseDomainsFile(filePath);
 
@@ -133,7 +133,7 @@
 
     it('should return empty array for file with only comments and whitespace', () => {
       const filePath = path.join(testDir, 'domains.txt');
       fs.writeFileSync(filePath, '# Comment 1\n\n# Comment 2\n   \n');
 
       const result = parseDomainsFile(filePath);
 
@@ -351,7 +351,7 @@
         .option('--env-all', 'Pass all env vars', false);
 
       // Parse empty args to get defaults
-      program.parse(['node', 'awf'], { from: 'user' });
+      program.parse([], { from: 'user' });
       const opts = program.opts();
 
       expect(opts.logLevel).toBe('info');
@@ -1396,13 +1396,22 @@
       expect(result.debugMessages[0]).toContain('Anthropic');
     });
 
-    it('should detect both keys', () => {
-      const result = validateApiProxyConfig(true, true, true);
+    it('should detect Copilot key', () => {
+      const result = validateApiProxyConfig(true, false, false, true);
       expect(result.enabled).toBe(true);
       expect(result.warnings).toEqual([]);
-      expect(result.debugMessages).toHaveLength(2);
+      expect(result.debugMessages).toHaveLength(1);
+      expect(result.debugMessages[0]).toContain('Copilot');
+    });
+
+    it('should detect all three keys', () => {
+      const result = validateApiProxyConfig(true, true, true, true);
+      expect(result.enabled).toBe(true);
+      expect(result.warnings).toEqual([]);
+      expect(result.debugMessages).toHaveLength(3);
       expect(result.debugMessages[0]).toContain('OpenAI');
       expect(result.debugMessages[1]).toContain('Anthropic');
+      expect(result.debugMessages[2]).toContain('Copilot');
     });
 
     it('should not warn when disabled even with keys', () => {

src/cli.ts

@@ -371,6 +371,19 @@ export interface FlagValidationResult {
   error?: string;
 }
 
+/**
+ * Checks if any rate limit options are set in the CLI options.
+ * Used to warn when rate limit flags are provided without --enable-api-proxy.
+ */
+export function hasRateLimitOptions(options: {
+  rateLimitRpm?: string;
+  rateLimitRph?: string;
+  rateLimitBytesPm?: string;
+  rateLimit?: boolean;
+}): boolean {
+  return !!(options.rateLimitRpm || options.rateLimitRph || options.rateLimitBytesPm || options.rateLimit === false);
+}
+
 /**
  * Validates that --skip-pull is not used with --build-local
  * @param skipPull - Whether --skip-pull flag was provided

tests/fixtures/awf-runner.ts

@@ -19,12 +19,12 @@ export interface AwfOptions {
   dnsServers?: string[]; // DNS servers to use (e.g., ['8.8.8.8', '2001:4860:4860::8888'])
   allowHostPorts?: string; // Ports or port ranges to allow for host access (e.g., '3000' or '3000-8000')
   enableApiProxy?: boolean; // Enable API proxy sidecar for LLM credential management
-  envAll?: boolean; // Pass all host environment variables to container (--env-all)
-  cliEnv?: Record<string, string>; // Explicit -e KEY=VALUE flags passed to AWF CLI
   rateLimitRpm?: number; // Requests per minute per provider
   rateLimitRph?: number; // Requests per hour per provider
   rateLimitBytesPm?: number; // Request bytes per minute per provider
   noRateLimit?: boolean; // Disable rate limiting
+  envAll?: boolean; // Pass all host environment variables to container (--env-all)
+  cliEnv?: Record<string, string>; // Explicit -e KEY=VALUE flags passed to AWF CLI
 }
 
 export interface AwfResult {
@@ -116,6 +116,20 @@ export class AwfRunner {
       args.push('--enable-api-proxy');
     }
 
+    // Add rate limit flags
+    if (options.rateLimitRpm !== undefined) {
+      args.push('--rate-limit-rpm', String(options.rateLimitRpm));
+    }
+    if (options.rateLimitRph !== undefined) {
+      args.push('--rate-limit-rph', String(options.rateLimitRph));
+    }
+    if (options.rateLimitBytesPm !== undefined) {
+      args.push('--rate-limit-bytes-pm', String(options.rateLimitBytesPm));
+    }
+    if (options.noRateLimit) {
+      args.push('--no-rate-limit');
+    }
+
     // Add --env-all flag
     if (options.envAll) {
       args.push('--env-all');
@@ -296,6 +310,20 @@ export class AwfRunner {
       args.push('--enable-api-proxy');
     }
 
+    // Add rate limit flags
+    if (options.rateLimitRpm !== undefined) {
+      args.push('--rate-limit-rpm', String(options.rateLimitRpm));
+    }
+    if (options.rateLimitRph !== undefined) {
+      args.push('--rate-limit-rph', String(options.rateLimitRph));
+    }
+    if (options.rateLimitBytesPm !== undefined) {
+      args.push('--rate-limit-bytes-pm', String(options.rateLimitBytesPm));
+    }
+    if (options.noRateLimit) {
+      args.push('--no-rate-limit');
+    }
+
     // Add --env-all flag
     if (options.envAll) {
       args.push('--env-all');

tests/integration/dns-servers.test.ts

@@ -10,7 +10,7 @@
 
 /// <reference path="../jest-custom-matchers.d.ts" />
 
-import { describe, test, expect, beforeAll, afterAll } from '@jest/globals';
+import { describe, test, expect, beforeAll, afterAll, beforeEach } from '@jest/globals';
 import { createRunner, AwfRunner } from '../fixtures/awf-runner';
 import { cleanup } from '../fixtures/cleanup';
 
@@ -113,3 +113,116 @@ describe('DNS Server Configuration', () => {
     expect(result.stdout.trim()).toMatch(/\d+\.\d+\.\d+\.\d+/);
   }, 120000);
 });
+
+describe('DNS Restriction Enforcement', () => {
+  let runner: AwfRunner;
+
+  beforeAll(async () => {
+    await cleanup(false);
+    runner = createRunner();
+  });
+
+  afterAll(async () => {
+    await cleanup(false);
+  });
+
+  // Clean up between each test to prevent container name conflicts
+  beforeEach(async () => {
+    await cleanup(false);
+  });
+
+  test('should block DNS queries to non-whitelisted servers', async () => {
+    // Only whitelist Google DNS (8.8.8.8) — Cloudflare (1.1.1.1) should be blocked
+    const result = await runner.runWithSudo(
+      'nslookup example.com 1.1.1.1',
+      {
+        allowDomains: ['example.com'],
+        dnsServers: ['8.8.8.8'],
+        logLevel: 'debug',
+        timeout: 60000,
+      }
+    );
+
+    // DNS query to non-whitelisted server should fail
+    expect(result).toFail();
+  }, 120000);
+
+  test('should allow DNS queries to whitelisted servers', async () => {
+    // Whitelist Google DNS (8.8.8.8) — queries to it should succeed
+    const result = await runner.runWithSudo(
+      'nslookup example.com 8.8.8.8',
+      {
+        allowDomains: ['example.com'],
+        dnsServers: ['8.8.8.8'],
+        logLevel: 'debug',
+        timeout: 60000,
+      }
+    );
+
+    expect(result).toSucceed();
+    expect(result.stdout).toContain('Address');
+  }, 120000);
+
+  test('should pass --dns-servers flag through to iptables configuration', async () => {
+    const result = await runner.runWithSudo(
+      'echo "dns-test"',
+      {
+        allowDomains: ['example.com'],
+        dnsServers: ['8.8.8.8'],
+        logLevel: 'debug',
+        timeout: 60000,
+      }
+    );
+
+    expect(result).toSucceed();
+    // Debug output should show the custom DNS server configuration
+    expect(result.stderr).toContain('8.8.8.8');
+  }, 120000);
+
+  test('should work with default DNS when --dns-servers is not specified', async () => {
+    // Without explicit dnsServers, default Google DNS (8.8.8.8, 8.8.4.4) should work
+    const result = await runner.runWithSudo(
+      'nslookup example.com',
+      {
+        allowDomains: ['example.com'],
+        logLevel: 'debug',
+        timeout: 60000,
+      }
+    );
+
+    expect(result).toSucceed();
+    expect(result.stdout).toContain('Address');
+  }, 120000);
+
+  test('should block DNS to non-default server when using defaults', async () => {
+    // With default DNS (8.8.8.8, 8.8.4.4), a query to a random DNS server
+    // like 208.67.222.222 (OpenDNS) should be blocked
+    const result = await runner.runWithSudo(
+      'nslookup example.com 208.67.222.222',
+      {
+        allowDomains: ['example.com'],
+        logLevel: 'debug',
+        timeout: 60000,
+      }
+    );
+
+    // DNS query to non-default server should fail
+    expect(result).toFail();
+  }, 120000);
+
+  test('should allow Cloudflare DNS when explicitly whitelisted', async () => {
+    // Whitelist Cloudflare DNS (1.1.1.1) — queries to it should succeed
+    const result = await runner.runWithSudo(
+      'nslookup example.com 1.1.1.1',
+      {
+        allowDomains: ['example.com'],
+        dnsServers: ['1.1.1.1'],
+        logLevel: 'debug',
+        timeout: 60000,
+      }
+    );
+
+    expect(result).toSucceed();
+    expect(result.stdout).toContain('Address');
+  }, 120000);
+});