diff --git a/src/cli.test.ts b/src/cli.test.ts index 7ad83e4b..d8790945 100644 --- a/src/cli.test.ts +++ b/src/cli.test.ts @@ -351,7 +351,7 @@ describe('cli', () => { .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 @@ describe('cli', () => { 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', () => { diff --git a/src/cli.ts b/src/cli.ts index b089e2f2..b06238e5 100644 --- a/src/cli.ts +++ b/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 diff --git a/tests/fixtures/awf-runner.ts b/tests/fixtures/awf-runner.ts index e137536d..aee732a0 100644 --- a/tests/fixtures/awf-runner.ts +++ b/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; // 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; // 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'); diff --git a/tests/integration/dns-servers.test.ts b/tests/integration/dns-servers.test.ts index cd471198..9054dafe 100644 --- a/tests/integration/dns-servers.test.ts +++ b/tests/integration/dns-servers.test.ts @@ -10,7 +10,7 @@ /// -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); +});