diff --git a/docs/acpremote.md b/docs/acpremote.md index 54d9b07..961134c 100644 --- a/docs/acpremote.md +++ b/docs/acpremote.md @@ -281,7 +281,7 @@ Current transport behavior is intentionally narrow: - one WebSocket text message carries one ACP JSON message - binary WebSocket frames are rejected - transport limits are configurable through `TransportOptions` -- bearer token auth is optional +- bearer token auth is optional; the token is read from the `Authorization: Bearer ` header or, for browser clients that cannot set headers, from a `bearer.` entry in `Sec-WebSocket-Protocol` (the server only echoes `acp.v1`, never the token) - metadata and health endpoints are served alongside the WebSocket transport ## Documented Remote-Host Flows diff --git a/packages/transports/acpremote/src/acpremote/auth.py b/packages/transports/acpremote/src/acpremote/auth.py index 8ec2d05..f9fb6ac 100644 --- a/packages/transports/acpremote/src/acpremote/auth.py +++ b/packages/transports/acpremote/src/acpremote/auth.py @@ -1,6 +1,7 @@ from __future__ import annotations as _annotations from collections.abc import Mapping +from hmac import compare_digest __all__ = ( "bearer_headers", @@ -9,6 +10,9 @@ ) +_BEARER_SUBPROTOCOL_PREFIX = "bearer." + + def bearer_headers(token: str | None) -> dict[str, str] | None: stripped = normalized_bearer_token(token) if stripped is None: @@ -29,7 +33,35 @@ def is_bearer_authorized( headers: Mapping[str, str], token: str | None, ) -> bool: + """Check the ``Authorization`` header, then ``bearer.`` WebSocket subprotocols. + + Browser WebSocket clients cannot set custom headers, so they carry the token as an extra + ``Sec-WebSocket-Protocol`` entry next to ``acp.v1``. + """ expected_token = normalized_bearer_token(token) if expected_token is None: return True - return headers.get("Authorization") == f"Bearer {expected_token}" + authorization = headers.get("Authorization") + if authorization is not None and compare_digest( + authorization.encode(), + f"Bearer {expected_token}".encode(), + ): + return True + return any( + compare_digest(candidate.encode(), expected_token.encode()) + for candidate in _subprotocol_bearer_tokens(headers) + ) + + +def _subprotocol_bearer_tokens(headers: Mapping[str, str]) -> list[str]: + try: + raw_subprotocols = headers.get("Sec-WebSocket-Protocol") + except LookupError: + return [] + if raw_subprotocols is None: + return [] + return [ + entry.removeprefix(_BEARER_SUBPROTOCOL_PREFIX) + for entry in (part.strip() for part in raw_subprotocols.split(",")) + if entry.startswith(_BEARER_SUBPROTOCOL_PREFIX) + ] diff --git a/tests/acpremote/test_helpers.py b/tests/acpremote/test_helpers.py index 8c33fd7..2a68157 100644 --- a/tests/acpremote/test_helpers.py +++ b/tests/acpremote/test_helpers.py @@ -21,6 +21,7 @@ SetSessionConfigOptionResponse, SetSessionModeResponse, ) +from acpremote import auth as auth_module from acpremote import client as client_module from acpremote import command as command_module from acpremote import proxy_agent as proxy_agent_module @@ -28,6 +29,7 @@ from acpremote.client import RemoteClientConnection from acpremote.config import TransportOptions from acpremote.metadata import ServerMetadata +from websockets.datastructures import Headers from websockets.exceptions import ConnectionClosedOK @@ -1198,3 +1200,11 @@ async def test_stream_helper_paths_cover_protocol_and_sender_edge_cases() -> Non ) with pytest.raises(asyncio.CancelledError): await cancelled_transport._sender_loop() + + +def test_bearer_authorization_rejects_duplicated_subprotocol_headers() -> None: + headers = Headers() + headers["Sec-WebSocket-Protocol"] = "bearer.secret-token" + headers["Sec-WebSocket-Protocol"] = "acp.v1" + + assert auth_module.is_bearer_authorized(headers, "secret-token") is False diff --git a/tests/acpremote/test_phase2.py b/tests/acpremote/test_phase2.py index c2cc35a..f4b5e73 100644 --- a/tests/acpremote/test_phase2.py +++ b/tests/acpremote/test_phase2.py @@ -285,6 +285,45 @@ async def test_phase2_websocket_path_and_bearer_auth_are_enforced() -> None: assert client.updates[0].update.content.text == "phase2 auth success" +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("subprotocols", "expected_status", "expected_subprotocol"), + [ + ([Subprotocol("acp.v1"), Subprotocol("bearer.secret-token")], None, "acp.v1"), + ([Subprotocol("bearer.secret-token")], None, None), + ([Subprotocol("acp.v1"), Subprotocol("bearer.wrong-token")], 401, None), + ([Subprotocol("acp.v1")], 401, None), + ], + ids=["acp-and-bearer", "bearer-only", "wrong-token", "no-token"], +) +async def test_phase2_websocket_accepts_bearer_token_from_subprotocol( + subprotocols: list[Subprotocol], + expected_status: int | None, + expected_subprotocol: str | None, +) -> None: + server = await serve_acp(cast("Agent", _EchoAgent()), bearer_token="secret-token") + assert server.sockets is not None + port = next(iter(server.sockets)).getsockname()[1] + try: + if expected_status is not None: + with pytest.raises(InvalidStatus) as rejected: + async with connect( + f"ws://127.0.0.1:{port}/acp/ws", + subprotocols=subprotocols, + ): + raise AssertionError("unreachable") # pragma: no cover + assert rejected.value.response.status_code == expected_status + else: + async with connect( + f"ws://127.0.0.1:{port}/acp/ws", + subprotocols=subprotocols, + ) as websocket: + assert websocket.subprotocol == expected_subprotocol + finally: + server.close() + await server.wait_closed() + + @pytest.mark.asyncio async def test_phase2_rejects_mixed_server_and_transport_options() -> None: with pytest.raises(ValueError, match="pass either options or server_options, not both"):