diff --git a/archivebox/core/asgi.py b/archivebox/core/asgi.py index 05f43287..50184274 100644 --- a/archivebox/core/asgi.py +++ b/archivebox/core/asgi.py @@ -13,5 +13,12 @@ from archivebox.config.django import setup_django setup_django(check_db=True) -# Standard Django ASGI application (no websockets/channels needed) -application = get_asgi_application() +django_application = get_asgi_application() + + +async def application(scope, receive, send): + if scope["type"] == "websocket": + from archivebox.opencode.views import websocket_view + + return await websocket_view(scope, receive, send) + return await django_application(scope, receive, send) diff --git a/archivebox/opencode/views.py b/archivebox/opencode/views.py index 46ccdfe7..43b09e5a 100644 --- a/archivebox/opencode/views.py +++ b/archivebox/opencode/views.py @@ -1,20 +1,43 @@ """Django adapter for the optional plugin; never import its runtime at startup.""" import logging +from io import BytesIO +from asgiref.sync import sync_to_async + +from django.contrib.auth.middleware import AuthenticationMiddleware from django.contrib.auth.views import redirect_to_login +from django.contrib.sessions.middleware import SessionMiddleware +from django.core.exceptions import PermissionDenied +from django.core.handlers.asgi import ASGIRequest +from django.db import close_old_connections from django.http import Http404, HttpResponse, HttpResponseForbidden, StreamingHttpResponse from django.template import engines +from django.urls import Resolver404, resolve from django.views.decorators.csrf import csrf_exempt from archivebox.config import CONSTANTS -from archivebox.config.common import get_config -from archivebox.core.routes_util import build_admin_url, get_api_base_url, get_base_url +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 _LOGGER = logging.getLogger(__name__) +def _runtime_settings(request, config): + from abx_plugins.plugins.opencode import runtime + + settings = runtime._settings(config, CONSTANTS.DATA_DIR) + route_config = request.__dict__.get("archivebox_config") + settings.update( + archivebox_base_url=get_base_url(request=request, config=route_config).rstrip("/"), + archivebox_admin_url=build_admin_url("/admin/", request=request, config=route_config).rstrip("/"), + archivebox_api_url=f"{get_api_base_url(request=request, config=route_config).rstrip('/')}/api/", + ) + return runtime, settings + + def _dispatch(request, path=None): try: config = dict(get_config().model_dump(mode="json")) @@ -25,15 +48,7 @@ def _dispatch(request, path=None): if not request.user.is_active or not request.user.is_superuser: return HttpResponseForbidden("Agent access requires a superuser account.") - from abx_plugins.plugins.opencode import runtime - - settings = runtime._settings(config, CONSTANTS.DATA_DIR) - route_config = request.__dict__.get("archivebox_config") - settings.update( - archivebox_base_url=get_base_url(request=request, config=route_config).rstrip("/"), - archivebox_admin_url=build_admin_url("/admin/", request=request, config=route_config).rstrip("/"), - archivebox_api_url=f"{get_api_base_url(request=request, config=route_config).rstrip('/')}/api/", - ) + runtime, settings = _runtime_settings(request, config) if path is None: from archivebox.core.admin_site import archivebox_admin @@ -74,3 +89,50 @@ def agent_view(request): @csrf_exempt def opencode_proxy_view(request, path=None): return _dispatch(request, path=path or "") + + +def _websocket_context(scope): + close_old_connections() + try: + match = resolve(scope["path"]) + if match.url_name != "opencode-proxy": + raise PermissionDenied + request = ASGIRequest( + {**scope, "type": "http", "method": "GET", "scheme": "https" if scope.get("scheme") == "wss" else "http"}, + BytesIO(), + ) + 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 + and not host_matches(request.get_host(), get_admin_host(config=route_config, request=request)) + ): + 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 + runtime, settings = _runtime_settings(request, config) + if not request.headers.get("Origin") or not runtime._origin_allowed("POST", request.get_host(), request.headers): + raise PermissionDenied + return runtime, settings, match.kwargs.get("path", "") + finally: + close_old_connections() + + +async def websocket_view(scope, receive, send): + if (await receive())["type"] != "websocket.connect": + return + try: + runtime, settings, path = await sync_to_async(_websocket_context)(scope) + await runtime.websocket_proxy(settings, path, scope.get("query_string", b""), scope.get("subprotocols", []), receive, send) + except (PermissionDenied, Resolver404, Http404): + await send({"type": "websocket.close", "code": 1008}) + except Exception: + _LOGGER.exception("Optional AI WebSocket failed") + await send({"type": "websocket.close", "code": 1011}) diff --git a/archivebox/tests/test_opencode_browser.py b/archivebox/tests/test_opencode_browser.py index 290646d8..5109fa2d 100644 --- a/archivebox/tests/test_opencode_browser.py +++ b/archivebox/tests/test_opencode_browser.py @@ -1,14 +1,18 @@ """Exercise the real agent wrapper with Chromium's native storage failures.""" +import asyncio import json import os +import re import subprocess import pytest +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 @@ -16,6 +20,93 @@ from .test_server_security_browser import browser_runtime as browser_runtime pytestmark = pytest.mark.django_db(transaction=True) +def test_agent_navigation_stays_inside_mount(agent_server, browser_runtime): + server_url, _, _ = agent_server + script = r""" +const assert = require('node:assert/strict'); +const puppeteer = require('puppeteer'); +const config = JSON.parse(require('node:fs').readFileSync(0, 'utf8')); +(async () => { + const browser = await puppeteer.launch({executablePath: config.chrome, headless: true, + args: ['--no-sandbox', '--disable-frame-rate-limit']}); + try { + const page = await browser.newPage(); + await page.setViewport({width: 1440, height: 900}); + await page.goto(config.url + '/admin/login/', {waitUntil: 'domcontentloaded'}); + await page.locator('#login-form input[name="username"]').fill('agent-browser-test'); + await page.locator('#login-form input[name="password"]').fill('test-password'); + await Promise.all([page.waitForNavigation({waitUntil: 'domcontentloaded'}), + page.locator('#login-form input[type="submit"]').click()]); + assert.equal((await page.goto(config.url + '/admin/agent', {waitUntil: 'domcontentloaded'})).status(), 200); + await page.locator('#opencode-agent-welcome-dismiss').click(); + const frame = await (await page.waitForSelector('iframe')).contentFrame(); + await frame.waitForSelector('a[href*="/session"]'); + const links = await frame.$$eval('a[href*="/session"]', nodes => nodes.map(node => node.getAttribute('href'))); + assert.ok(links.length, 'OpenCode must expose session navigation'); + for (const href of links) assert.ok(href.startsWith('/admin/agent/opencode/'), href); + await frame.locator('::-p-aria(New session[role="button"])').click(); + await frame.waitForSelector('[contenteditable="true"]'); + let navigation; + for (const link of await frame.$$('a[href*="/session"]')) { + if (await link.isVisible() && await link.evaluate(node => node.href) !== frame.url()) { + navigation = link; + break; + } + } + assert.ok(navigation, 'A visible link must navigate to another session route'); + await navigation.click(); + await frame.waitForSelector('[contenteditable="true"]'); + assert.ok(new URL(frame.url()).pathname.startsWith('/admin/agent/opencode/'), frame.url()); + const sessionUrl = frame.url(); + assert.equal((await frame.goto(sessionUrl, {waitUntil: 'domcontentloaded'})).status(), 200); + await frame.waitForSelector('[contenteditable="true"]'); + assert.ok(!(await frame.$eval('body', node => node.innerText)).includes('Something went wrong')); + // Exercise the actual public PTY API and native browser WebSocket. No + // intercepted traffic or replacement server: this runs a real shell. + const terminal = await frame.evaluate(async () => { + const base = location.origin + '/admin/agent/opencode'; + const created = await fetch(base + '/pty', { + method: 'POST', headers: {'Content-Type': 'application/json'}, + body: JSON.stringify({command: '/bin/sh', args: []}), + }); + if (!created.ok) throw new Error('PTY create: ' + created.status); + const pty = await created.json(); + try { + return await new Promise((resolve, reject) => { + const url = new URL(base + '/pty/' + pty.id + '/connect'); + url.protocol = url.protocol === 'https:' ? 'wss:' : 'ws:'; + const socket = new WebSocket(url); + const timer = setTimeout(() => { socket.close(); reject(new Error('PTY output timed out')); }, 10000); + let output = ''; + socket.onopen = () => socket.send("printf 'ABX_%s\\n' TERMINAL_OK\n"); + socket.onmessage = async event => { + output += typeof event.data === 'string' ? event.data : await event.data.text(); + if (output.includes('ABX_TERMINAL_OK')) { + clearTimeout(timer); socket.close(); resolve(output); + } + }; + socket.onerror = () => { clearTimeout(timer); reject(new Error('PTY WebSocket failed')); }; + }); + } finally { await fetch(base + '/pty/' + pty.id, {method: 'DELETE'}); } + }); + assert.ok(terminal.includes('ABX_TERMINAL_OK'), terminal); + assert.equal((await page.goto(config.url + '/add/', {waitUntil: 'domcontentloaded'})).status(), 200); + console.log('AGENT_NAVIGATION_OK'); + } finally { await browser.close(); } +})().catch(error => { console.error(error); process.exitCode = 1; }); +""" + result = subprocess.run( + [str(browser_runtime["node_binary"]), "-e", script], + input=json.dumps({"chrome": str(browser_runtime["chrome_binary"]), "url": server_url}), + env={**os.environ, "NODE_PATH": browser_runtime["node_path"]}, + text=True, + capture_output=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr or result.stdout + assert "AGENT_NAVIGATION_OK" in result.stdout + + @pytest.fixture def agent_server(installed_opencode, browser_runtime): port = get_free_port() @@ -30,7 +121,7 @@ def agent_server(installed_opencode, browser_runtime): [ "shell", "-c", - "from django.contrib.auth import get_user_model; get_user_model().objects.create_superuser(username='agent-browser-test', password='test-password')", + "from django.contrib.auth import get_user_model; User = get_user_model(); User.objects.create_superuser(username='agent-browser-test', password='test-password'); User.objects.create_user(username='agent-regular-test', password='test-password', is_staff=True)", ], cwd=config.data_dir, env=config.env, @@ -38,13 +129,132 @@ 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 + admin_cookie = _login_cookie(server_url) + regular_cookie = _login_cookie(server_url, "agent-regular-test") + + async def check_denials(): + for path, cookie, origin in ( + ("/admin/agent/opencode/pty/unknown/connect", "", server_url), + ("/admin/agent/opencode/pty/unknown/connect", regular_cookie, server_url), + ("/admin/agent/opencode/pty/unknown/connect", admin_cookie, "https://untrusted.example"), + ("/admin/agent/opencode/pty/unknown/connect", admin_cookie, None), + ("/health/", admin_cookie, server_url), + ): + with pytest.raises(InvalidStatus) as error: + async with connect( + server_url.replace("http", "ws", 1) + path, + origin=origin, + additional_headers={"Cookie": cookie}, + proxy=None, + ): + pytest.fail("Unauthorized WebSocket was accepted") + assert error.value.response.status_code == 403 + + asyncio.run(check_denials()) + 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'); diff --git a/pyproject.toml b/pyproject.toml index 9ad18f9f..a3e80449 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,7 +42,7 @@ dependencies = [ "setuptools>=74.1.0", # for: django 5 on python >=3.12, distutils is no longer in stdlib but django 5.1 expects distutils (TODO: check if this can be removed eventually) "django>=6.1", "psycopg[binary]>=3.2", # for: PostgreSQL database backend (ARCHIVEBOX_DATABASE_ENGINE=postgres) - "daphne>=4.2.1", # ASGI server for Django (no channels needed - websockets not used) + "daphne>=4.2.1", # ASGI server for Django and optional plugin WebSockets (no channels needed) "django-ninja>=1.5.1", "django-extensions>=3.2.3", "django-signal-webhooks>=0.3.0",