From 07da459bb1dfb9d62e381c956115c5406c848aa8 Mon Sep 17 00:00:00 2001 From: Nick Sweeting Date: Thu, 3 Sep 2026 17:27:45 -0700 Subject: [PATCH] Apply control-plane policy and proxy authentication to agent sockets --- archivebox/opencode/views.py | 4 + archivebox/tests/test_opencode_browser.py | 122 +++++++++++++++++----- 2 files changed, 101 insertions(+), 25 deletions(-) diff --git a/archivebox/opencode/views.py b/archivebox/opencode/views.py index bce3e4a6..43b09e5a 100644 --- a/archivebox/opencode/views.py +++ b/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 diff --git a/archivebox/tests/test_opencode_browser.py b/archivebox/tests/test_opencode_browser.py index 7ca4691b..5109fa2d 100644 --- a/archivebox/tests/test_opencode_browser.py +++ b/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');