tests/test_controller_auth.py

"""Browser credential lifecycle and authenticated controller compatibility."""

from __future__ import annotations

import hashlib
import hmac
import http.client
import json
import subprocess
import tempfile
import threading
import unittest
from pathlib import Path
from unittest import mock

from test_mcp_server import CONTROLLER_PATH, load_server


class ControllerAuthTests(unittest.TestCase):
    def test_legacy_controller_is_verified_but_not_reusable(self) -> None:
        server = load_server()
        token = "synthetic-control-" + "x" * 48

        def reply(method, url):
            challenge = server.urllib.parse.parse_qs(server.urllib.parse.urlsplit(url).query)[
                "challenge"
            ][0]
            return 200, json.dumps(
                {
                    "server": server.SERVER_NAME,
                    "protocol_version": 1,
                    "proof": hmac.new(
                        token.encode(), challenge.encode(), hashlib.sha256
                    ).hexdigest(),
                }
            ).encode()

        with (
            mock.patch.object(
                server, "_read_controller_runtime", return_value={"control_token": token}
            ),
            mock.patch.object(server, "_http_request", side_effect=reply),
        ):
            status, _ = server._probe_controller(43117)
        self.assertEqual(status, "legacy")

    def test_browser_fragment_refresh_invalid_link_and_api_headers(self) -> None:
        harness = r"""
const fs = require('node:fs'), vm = require('node:vm'), assert = require('node:assert/strict');
const file = process.argv[1], source = fs.readFileSync(file, 'utf8');
const script = source.split('<script>')[1].split('</script>')[0];
const saved = new Map(), requests = [];
function load(hash, blockedStorage = false) {
  const app = {innerHTML:'',addEventListener() {}};
  const location = {hash,pathname:'/',search:''};
  const context = vm.createContext({
    console, location,
    window:{addEventListener(){}},
    history:{replaceState(a,b,url){assert.equal(url,'/');location.hash='';}},
    sessionStorage:{getItem:k=>saved.get(k),setItem(k,v){if(blockedStorage)throw Error('blocked');saved.set(k,v);},removeItem:k=>saved.delete(k)},
    localStorage:{getItem:()=>null,setItem(){},removeItem(){}},
    document:{getElementById:()=>app,createElement:()=>({click(){},remove(){}}),body:{append(){}}},
    URL:{createObjectURL:()=> 'blob:fixture',revokeObjectURL(){}},
    fetch:async (url, options)=>{requests.push({url,options});return {ok:true,json:async()=>({structuredContent:{ok:true}}),blob:async()=>({})};},
  });
  vm.runInContext(fs.readFileSync(require('node:path').join(require('node:path').dirname(file),'controller-ranges.js'),'utf8'),context);
  vm.runInContext(script.replace('      initialize();', 'render = () => {}; window.test = {accessToken, callTool, heartbeatController, downloadReport, state};'), context);
  assert.equal(location.hash,'');
  return context.window.test;
}
(async()=>{
 const token='a'.repeat(64), fresh='b'.repeat(64);
 let api=load('#token='+token); assert.equal(api.accessToken,token);
 api=load(''); assert.equal(api.accessToken,token);
 api.state.runId='fixture'; await api.callTool('get_state'); await api.heartbeatController(); await api.downloadReport('json');
 assert.deepEqual(requests.map(r=>r.url),['/api/call','/api/heartbeat','/api/download']);
 for(const r of requests)assert.equal(r.options.headers['X-Codex-Replay-Access'],token);
 assert.equal(load('#token='+fresh).accessToken,fresh);
 assert.equal(load('#token=malformed').accessToken,''); assert.equal(load('').accessToken,'');
 assert.equal(load('#token='+token,true).accessToken,token);
})().catch(e=>{console.error(e);process.exitCode=1;});
"""
        subprocess.run(
            ["node", "-e", harness, str(CONTROLLER_PATH)],
            check=True,
            capture_output=True,
            text=True,
            timeout=15,
        )

    def test_loopback_controller_requires_browser_authentication(self) -> None:
        server = load_server()
        control_token = "control-" + ("x" * 48)
        httpd = server._ControllerHTTPServer(("127.0.0.1", 0), control_token)
        thread = threading.Thread(target=httpd.serve_forever, daemon=True)
        thread.start()
        port = int(httpd.server_address[1])

        def request(
            method: str,
            path: str,
            *,
            body: str | None = None,
            headers: dict[str, str] | None = None,
        ) -> tuple[int, dict[str, str], bytes]:
            connection = http.client.HTTPConnection("127.0.0.1", port, timeout=5)
            try:
                connection.request(
                    method,
                    path,
                    body=body,
                    headers={
                        server.CONTROLLER_BROWSER_HEADER: httpd.browser_token,
                        **(headers or {}),
                    },
                )
                response = connection.getresponse()
                return response.status, dict(response.getheaders()), response.read()
            finally:
                connection.close()

        try:
            with (
                mock.patch.object(server, "_call_tool") as dispatch,
                mock.patch.object(httpd, "touch_heartbeat") as heartbeat,
                mock.patch.object(server, "_safe_run_directory") as read_run,
            ):
                for route in ("/api/call", "/api/download", "/api/heartbeat"):
                    for token in ("", "wrong", control_token, "x" * 64):
                        status, _, _ = request(
                            "POST",
                            route,
                            body="{}",
                            headers={
                                "Origin": httpd.origin,
                                "Content-Type": "application/json",
                                server.CONTROLLER_BROWSER_HEADER: token,
                            },
                        )
                        self.assertEqual(status, 401, (route, token))
                dispatch.assert_not_called()
                heartbeat.assert_not_called()
                read_run.assert_not_called()
            status, _, _ = request(
                "POST",
                "/api/call",
                body="{}",
                headers={
                    "Host": "attacker.invalid",
                    "Origin": httpd.origin,
                    "Content-Type": "application/json",
                },
            )
            self.assertEqual(status, 421)
            challenge = "test-challenge"
            status, _, body = request("GET", f"/health?challenge={challenge}")
            self.assertEqual(status, 200)
            health = json.loads(body)
            self.assertEqual(health["server"], "codex-replay")
            self.assertNotIn(httpd.browser_token.encode(), body)
            self.assertNotIn(control_token.encode(), body)
            self.assertEqual(
                health["proof"],
                hmac.new(control_token.encode(), challenge.encode(), hashlib.sha256).hexdigest(),
            )

            status, _, _ = request("GET", "/auth?token=unused")
            self.assertEqual(status, 404)
            status, _, _ = request("POST", "/api/launch", body="{}")
            self.assertEqual(status, 404)

            status, _, body = request(
                "POST",
                "/api/heartbeat",
                body="{}",
                headers={
                    "Content-Type": "application/json",
                    "Origin": httpd.origin,
                },
            )
            self.assertEqual(status, 200)
            heartbeat = json.loads(body)
            self.assertTrue(heartbeat["ok"])
            self.assertEqual(heartbeat["controller_session_id"], httpd.controller_session_id)
            self.assertGreater(heartbeat["heartbeat_interval_seconds"], 0)

            status, _, _ = request(
                "POST",
                "/api/heartbeat",
                body="{}",
                headers={
                    "Content-Type": "application/json",
                    "Origin": "https://example.com",
                },
            )
            self.assertEqual(status, 403)

            status, _, body = request("GET", "/")
            self.assertEqual(status, 200)
            self.assertIn(b"codex-replay.controller-draft.v7", body)
            self.assertNotIn(httpd.browser_token.encode(), body)
            self.assertNotIn(b"window.openai", body)
            self.assertNotIn(b"controller-session", body)
            status, headers, body = request("GET", "/controller.css")
            self.assertEqual(status, 200)
            self.assertEqual(headers["Content-Type"], "text/css; charset=utf-8")
            self.assertIn(b":root", body)
            status, headers, body = request("GET", "/controller-ranges.js")
            self.assertEqual(status, 200)
            self.assertEqual(headers["Content-Type"], "text/javascript; charset=utf-8")
            self.assertEqual(body, server.APP_RANGES.read_bytes())
            self.assertIn("script-src 'self'", headers["Content-Security-Policy"])

            result = {
                "content": [{"type": "text", "text": "Ready."}],
                "structuredContent": {"state": {"models": []}},
            }
            with mock.patch.object(server, "_call_tool", return_value=result):
                status, _, body = request(
                    "POST",
                    "/api/call",
                    body=json.dumps({"name": "get_state", "arguments": {}}),
                    headers={
                        "Content-Type": "application/json",
                        "Origin": httpd.origin,
                    },
                )
                self.assertEqual(status, 200)
                self.assertEqual(json.loads(body), result)

                status, _, _ = request(
                    "POST",
                    "/api/call",
                    body=json.dumps({"name": "get_state", "arguments": {}}),
                    headers={
                        "Content-Type": "application/json",
                        "Origin": "https://example.com",
                    },
                )
                self.assertEqual(status, 403)

                status, _, _ = request(
                    "POST",
                    "/api/call",
                    body=json.dumps({"name": "get_state", "arguments": {}}),
                    headers={"Content-Type": "application/json"},
                )
                self.assertEqual(status, 403)

                status, _, _ = request(
                    "POST",
                    "/api/call",
                    body=json.dumps({"name": "open_controller", "arguments": {}}),
                    headers={
                        "Content-Type": "application/json",
                        "Origin": httpd.origin,
                    },
                )
                self.assertEqual(status, 400)
        finally:
            httpd.shutdown()
            thread.join(timeout=5)
            httpd.server_close()

    def test_authenticated_same_origin_report_downloads(self) -> None:
        server = load_server()
        with tempfile.TemporaryDirectory() as temporary:
            run_root = Path(temporary).resolve()
            run_directory = run_root / "run-1"
            run_directory.mkdir()
            artifacts = {
                "json": (
                    b'{"status":"completed","original_request":"'
                    + b"x" * (19 * 1024 * 1024)
                    + b'"}\n',
                    "application/json; charset=utf-8",
                ),
                "html": (
                    b"<!doctype html><title>Replay report</title>" + b"x" * (19 * 1024 * 1024),
                    "text/html; charset=utf-8",
                ),
            }
            for artifact_format, (body, _) in artifacts.items():
                (run_directory / f"report.{artifact_format}").write_bytes(body)

            with mock.patch.object(server, "RUN_ROOT", run_root):
                httpd = server._ControllerHTTPServer(
                    ("127.0.0.1", 0),
                    "control-" + ("x" * 48),
                )
                thread = threading.Thread(target=httpd.serve_forever, daemon=True)
                thread.start()
                port = int(httpd.server_address[1])

                def request(
                    artifact_format: str,
                    *,
                    origin: str | None = None,
                ) -> tuple[int, dict[str, str], bytes]:
                    connection = http.client.HTTPConnection("127.0.0.1", port, timeout=5)
                    headers = {
                        "Content-Type": "application/json",
                        "Origin": httpd.origin if origin is None else origin,
                        server.CONTROLLER_BROWSER_HEADER: httpd.browser_token,
                    }
                    try:
                        connection.request(
                            "POST",
                            "/api/download",
                            body=json.dumps({"run_id": "run-1", "format": artifact_format}),
                            headers=headers,
                        )
                        response = connection.getresponse()
                        return response.status, dict(response.getheaders()), response.read()
                    finally:
                        connection.close()

                try:
                    for artifact_format, (expected_body, content_type) in artifacts.items():
                        status, headers, body = request(artifact_format)
                        self.assertEqual(status, 200)
                        self.assertEqual(body, expected_body)
                        self.assertEqual(headers["Content-Type"], content_type)
                        self.assertEqual(
                            headers["Content-Disposition"],
                            f'attachment; filename="codex-replay-run-1-report.{artifact_format}"',
                        )

                    status, _, _ = request("html", origin="https://example.com")
                    self.assertEqual(status, 403)
                finally:
                    httpd.shutdown()
                    thread.join(timeout=5)
                    httpd.server_close()