|
| 1 | +"""Toggle a running container's internet access via iptables in its network namespace. |
| 2 | +
|
| 3 | +Ported from RevEngBench (reveng/utils/environment_internet_control.py), Linux-only path. |
| 4 | +Used to disable internet DURING the agent's editing turn (so it can't look up solutions), |
| 5 | +while leaving it on for setup (git fetch of the starting codebase) and the post-round push. |
| 6 | +Loopback is always kept up, so localhost gameplay (bot servers on 0.0.0.0:PORT) is unaffected. |
| 7 | +DNS is left reachable so container tools fail fast instead of hanging on resolution. |
| 8 | +
|
| 9 | +Host-side (nsenter + sudo iptables) so an in-container root agent cannot undo it. Composes with |
| 10 | +Player._isolate_git — internet-off stops re-fetching, git-strip hides the already-cloned branches. |
| 11 | +""" |
| 12 | + |
| 13 | +import os |
| 14 | +import subprocess |
| 15 | +import threading |
| 16 | + |
| 17 | +from minisweagent.environments.docker import DockerEnvironment |
| 18 | + |
| 19 | +from codeclash.utils.log import get_logger |
| 20 | + |
| 21 | +logger = get_logger(__name__) |
| 22 | + |
| 23 | +_dns_lock = threading.Lock() |
| 24 | +_container_dns: dict[str, set[str]] = {} |
| 25 | + |
| 26 | + |
| 27 | +def _docker_pid(env: DockerEnvironment) -> str: |
| 28 | + return subprocess.run( |
| 29 | + [env.config.executable, "inspect", "--format", "{{.State.Pid}}", env.container_id], |
| 30 | + capture_output=True, |
| 31 | + text=True, |
| 32 | + check=True, |
| 33 | + ).stdout.strip() |
| 34 | + |
| 35 | + |
| 36 | +def _iptables(env: DockerEnvironment, rule_args: list[str]) -> None: |
| 37 | + """Run one iptables command in the container's net namespace (sudo nsenter on Linux).""" |
| 38 | + pid = _docker_pid(env) |
| 39 | + prefix = [] if os.getuid() == 0 else ["sudo", "-n"] |
| 40 | + subprocess.run([*prefix, "nsenter", "-t", pid, "-n", "iptables", *rule_args], check=True) |
| 41 | + |
| 42 | + |
| 43 | +def _dns_servers(env: DockerEnvironment) -> set[str]: |
| 44 | + r = env.execute({"command": "grep '^nameserver' /etc/resolv.conf | awk '{print $2}'"}) |
| 45 | + return {ip.strip() for ip in r["output"].splitlines() if ip.strip()} |
| 46 | + |
| 47 | + |
| 48 | +def turn_off_internet(env: DockerEnvironment) -> None: |
| 49 | + """Drop all inbound/outbound traffic except loopback and DNS.""" |
| 50 | + dns = _dns_servers(env) |
| 51 | + with _dns_lock: |
| 52 | + _container_dns[env.container_id] = dns |
| 53 | + for chain in ("OUTPUT", "INPUT"): |
| 54 | + _iptables(env, ["-I", chain, "-j", "DROP"]) |
| 55 | + _iptables(env, ["-I", "OUTPUT", "-o", "lo", "-j", "ACCEPT"]) |
| 56 | + _iptables(env, ["-I", "INPUT", "-i", "lo", "-j", "ACCEPT"]) |
| 57 | + for ip in sorted(dns): |
| 58 | + _iptables(env, ["-I", "OUTPUT", "-d", ip, "-p", "udp", "--dport", "53", "-j", "ACCEPT"]) |
| 59 | + _iptables(env, ["-I", "OUTPUT", "-d", ip, "-p", "tcp", "--dport", "53", "-j", "ACCEPT"]) |
| 60 | + _iptables(env, ["-I", "INPUT", "-s", ip, "-p", "udp", "--sport", "53", "-j", "ACCEPT"]) |
| 61 | + _iptables(env, ["-I", "INPUT", "-s", ip, "-p", "tcp", "--sport", "53", "-j", "ACCEPT"]) |
| 62 | + logger.info("Internet disabled for container %s (loopback+DNS kept)", env.container_id[:12]) |
| 63 | + |
| 64 | + |
| 65 | +def turn_on_internet(env: DockerEnvironment) -> None: |
| 66 | + """Remove the drop rules inserted by turn_off_internet (restores full connectivity).""" |
| 67 | + with _dns_lock: |
| 68 | + dns = _container_dns.pop(env.container_id, set()) |
| 69 | + for ip in sorted(dns): |
| 70 | + _iptables(env, ["-D", "OUTPUT", "-d", ip, "-p", "udp", "--dport", "53", "-j", "ACCEPT"]) |
| 71 | + _iptables(env, ["-D", "OUTPUT", "-d", ip, "-p", "tcp", "--dport", "53", "-j", "ACCEPT"]) |
| 72 | + _iptables(env, ["-D", "INPUT", "-s", ip, "-p", "udp", "--sport", "53", "-j", "ACCEPT"]) |
| 73 | + _iptables(env, ["-D", "INPUT", "-s", ip, "-p", "tcp", "--sport", "53", "-j", "ACCEPT"]) |
| 74 | + _iptables(env, ["-D", "OUTPUT", "-o", "lo", "-j", "ACCEPT"]) |
| 75 | + _iptables(env, ["-D", "INPUT", "-i", "lo", "-j", "ACCEPT"]) |
| 76 | + for chain in ("OUTPUT", "INPUT"): |
| 77 | + _iptables(env, ["-D", chain, "-j", "DROP"]) |
| 78 | + logger.info("Internet restored for container %s", env.container_id[:12]) |
0 commit comments