Apply control-plane policy and proxy authentication to agent sockets

archivebox/opencode/views.py

@@ -18,6 +18,7 @@ from django.views.decorators.csrf import csrf_exempt
 
 from archivebox.config import CONSTANTS
 from archivebox.config.common import get_config, get_request_config
+from archivebox.core.middleware import ReverseProxyAuthMiddleware
 from archivebox.core.routes_util import build_admin_url, get_admin_host, get_api_base_url, get_base_url, host_matches
 from archivebox.plugins.discovery import get_plugin_template
 
@@ -102,6 +103,8 @@ def _websocket_context(scope):
         )
         route_config = get_request_config(request, resolve_plugins=False)
         request.archivebox_config = route_config
+        if not route_config.CONTROL_PLANE_ENABLED:
+            raise PermissionDenied
         if (
             route_config.USES_SUBDOMAIN_ROUTING
             and route_config.BASE_URL
@@ -110,6 +113,7 @@ def _websocket_context(scope):
             raise PermissionDenied
         SessionMiddleware(_dispatch).process_request(request)
         AuthenticationMiddleware(_dispatch).process_request(request)
+        ReverseProxyAuthMiddleware(_dispatch).process_request(request)
         config = get_config().model_dump(mode="json")
         if not config.get("OPENCODE_ENABLED") or not request.user.is_active or not request.user.is_superuser:
             raise PermissionDenied

archivebox/tests/test_opencode_browser.py

@@ -12,6 +12,7 @@ import requests
 from .conftest import get_free_port, run_archivebox_cmd, start_archivebox_server, stop_archivebox_process
 from .test_opencode_agent import _set_archivebox_config
 from .test_opencode_agent import installed_opencode as installed_opencode
+from .test_opencode_agent import live_opencode as live_opencode
 from .test_opencode_agent import opencode_archive_config as opencode_archive_config
 from .test_server_security_browser import browser_runtime as browser_runtime
 
@@ -20,7 +21,7 @@ pytestmark = pytest.mark.django_db(transaction=True)
 
 
 def test_agent_navigation_stays_inside_mount(agent_server, browser_runtime):
-    server_url, _ = agent_server
+    server_url, _, _ = agent_server
     script = r"""
 const assert = require('node:assert/strict');
 const puppeteer = require('puppeteer');
@@ -128,36 +129,37 @@ def agent_server(installed_opencode, browser_runtime):
     assert user.returncode == 0, user.stderr or user.stdout
     process = start_archivebox_server(config.data_dir, port=port, env=config.env, log_name="agent-browser-server.log")
     try:
-        yield url, config.data_dir
+        yield url, config.data_dir, process
     finally:
-        stop_archivebox_process(process)
+        if process.poll() is None:
+            stop_archivebox_process(process)
+
+
+def _login_cookie(server_url, username="agent-browser-test"):
+    session = requests.Session()
+    login_url = server_url + "/admin/login/"
+    page = session.get(login_url, timeout=10)
+    assert page.status_code == 200
+    token = re.search(r'name="csrfmiddlewaretoken" value="([^"]+)"', page.text)
+    assert token is not None
+    response = session.post(
+        login_url,
+        data={"username": username, "password": "test-password", "csrfmiddlewaretoken": token[1]},
+        headers={"Referer": login_url},
+        allow_redirects=False,
+        timeout=10,
+    )
+    assert response.status_code == 302
+    return "; ".join(f"{key}={value}" for key, value in session.cookies.items())
 
 
 def test_agent_websocket_rejects_unauthorized_access(agent_server):
     from websockets.asyncio.client import connect
     from websockets.exceptions import InvalidStatus
 
-    server_url, _ = agent_server
-
-    def login(username):
-        session = requests.Session()
-        login_url = server_url + "/admin/login/"
-        page = session.get(login_url, timeout=10)
-        assert page.status_code == 200
-        token = re.search(r'name="csrfmiddlewaretoken" value="([^"]+)"', page.text)
-        assert token is not None
-        response = session.post(
-            login_url,
-            data={"username": username, "password": "test-password", "csrfmiddlewaretoken": token[1]},
-            headers={"Referer": login_url},
-            allow_redirects=False,
-            timeout=10,
-        )
-        assert response.status_code == 302
-        return "; ".join(f"{key}={value}" for key, value in session.cookies.items())
-
-    admin_cookie = login("agent-browser-test")
-    regular_cookie = login("agent-regular-test")
+    server_url, _, _ = agent_server
+    admin_cookie = _login_cookie(server_url)
+    regular_cookie = _login_cookie(server_url, "agent-regular-test")
 
     async def check_denials():
         for path, cookie, origin in (
@@ -181,8 +183,78 @@ def test_agent_websocket_rejects_unauthorized_access(agent_server):
     assert requests.get(server_url + "/health/", timeout=10).status_code == 200
 
 
+@pytest.mark.parametrize(
+    "mode,proxy_whitelist,use_cookie,http_status",
+    [
+        ("unsafe-onedomain-noadmin", "", True, 403),
+        ("safe-onedomain-nojsreplay", "127.0.0.1/32", False, 200),
+        ("safe-onedomain-nojsreplay", "192.0.2.0/24", False, 302),
+    ],
+)
+def test_agent_websocket_matches_http_security_policy(agent_server, live_opencode, mode, proxy_whitelist, use_cookie, http_status):
+    from urllib.parse import urlsplit
+    from websockets.asyncio.client import connect
+    from websockets.exceptions import InvalidStatus
+
+    server_url, data_dir, process = agent_server
+    cookie = _login_cookie(server_url)
+    created = requests.post(
+        server_url + "/admin/agent/opencode/pty",
+        headers={"Cookie": cookie, "Origin": server_url},
+        json={"command": "/bin/sh", "args": []},
+        timeout=10,
+    )
+    assert created.status_code == 200
+    pty_id = created.json()["id"]
+    stop_archivebox_process(process)
+    _set_archivebox_config(
+        data_dir,
+        f"SERVER_SECURITY_MODE={mode}",
+        f"REVERSE_PROXY_WHITELIST={proxy_whitelist}",
+        "REVERSE_PROXY_USER_HEADER=Remote-User",
+    )
+    restarted = start_archivebox_server(
+        data_dir,
+        port=urlsplit(server_url).port,
+        env=live_opencode.config.env,
+        log_name="agent-security-server.log",
+    )
+    headers = {"Cookie": cookie} if use_cookie else {"Remote-User": "agent-browser-test"}
+    try:
+        response = requests.get(server_url + "/admin/agent/opencode/global/health", headers=headers, allow_redirects=False, timeout=10)
+        assert response.status_code == http_status
+
+        async def check_socket():
+            connection = connect(
+                server_url.replace("http", "ws", 1) + f"/admin/agent/opencode/pty/{pty_id}/connect",
+                origin=server_url,
+                additional_headers=headers,
+                proxy=None,
+            )
+            if http_status != 200:
+                with pytest.raises(InvalidStatus) as error:
+                    async with connection:
+                        pytest.fail("Disabled control plane or untrusted proxy accepted a WebSocket")
+                assert error.value.response.status_code == 403
+                return
+            async with connection as socket, asyncio.timeout(10):
+                await socket.send("printf 'ABX_%s\\n' PROXY_OK\n")
+                output = ""
+                while "ABX_PROXY_OK" not in output:
+                    chunk = await socket.recv()
+                    output += chunk.decode() if isinstance(chunk, bytes) else chunk
+                assert "ABX_PROXY_OK" in output
+
+        asyncio.run(check_socket())
+        assert requests.get(server_url + "/health/", timeout=10).status_code == 200
+    finally:
+        deleted = requests.delete(live_opencode.settings["origin"] + f"/pty/{pty_id}", timeout=10)
+        assert deleted.status_code == 200
+        stop_archivebox_process(restarted)
+
+
 def test_agent_preserves_projects_and_survives_storage_failure(agent_server, browser_runtime):
-    server_url, data_dir = agent_server
+    server_url, data_dir, _ = agent_server
     script = r"""
 const assert = require('node:assert/strict');
 const puppeteer = require('puppeteer');