fix: properly authenticate terminal route (#9098)

Author: mscolnick

marimo/_server/api/endpoints/terminal.py

@@ -14,7 +14,9 @@
 from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState
 
 from marimo import _loggers
+from marimo._server.api.auth import validate_auth
 from marimo._server.api.deps import AppState
+from marimo._server.codes import WebSocketCodes
 from marimo._server.router import APIRouter
 from marimo._session.model import SessionMode
 from marimo._utils.platform import is_pyodide, is_windows
@@ -353,6 +355,13 @@ def supports_terminal() -> bool:
 @router.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket) -> None:
     app_state = AppState(websocket)
+
+    if app_state.enable_auth and not validate_auth(websocket):
+        await websocket.close(
+            WebSocketCodes.UNAUTHORIZED, "MARIMO_UNAUTHORIZED"
+        )
+        return
+
     if app_state.mode != SessionMode.EDIT:
         await websocket.close(
             code=1008, reason="Terminal only available in edit mode"

tests/_server/api/endpoints/test_terminal.py

@@ -24,6 +24,8 @@
 from marimo._session.model import SessionMode
 from tests._server.mocks import get_session_manager
 
+TERMINAL_WS_URL = "/terminal/ws?access_token=fake-token"
+
 if TYPE_CHECKING:
     from starlette.testclient import TestClient
 
@@ -33,7 +35,7 @@
 
 @pytest.mark.skipif(is_windows, reason="Skip on Windows")
 def test_terminal_ws(client: TestClient) -> None:
-    with client.websocket_connect("/terminal/ws") as websocket:
+    with client.websocket_connect(TERMINAL_WS_URL) as websocket:
         # Send echo message
         websocket.send_text("echo hello")
         data = websocket.receive_text()
@@ -44,11 +46,27 @@ def test_terminal_ws_not_allowed_in_run(client: TestClient) -> None:
     session_manager: SessionManager = get_session_manager(client)
     session_manager.mode = SessionMode.RUN
     with pytest.raises(WebSocketDisconnect):
-        with client.websocket_connect("/terminal/ws") as websocket:
+        with client.websocket_connect(TERMINAL_WS_URL) as websocket:
             websocket.send_text("echo hello")
     session_manager.mode = SessionMode.EDIT
 
 
+def test_terminal_ws_unauthorized(client: TestClient) -> None:
+    """Test terminal websocket rejects unauthenticated connections."""
+    with pytest.raises(WebSocketDisconnect) as exc_info:
+        with client.websocket_connect("/terminal/ws"):
+            pass
+    assert exc_info.value.code == 3000
+
+
+def test_terminal_ws_wrong_token(client: TestClient) -> None:
+    """Test terminal websocket rejects wrong token."""
+    with pytest.raises(WebSocketDisconnect) as exc_info:
+        with client.websocket_connect("/terminal/ws?access_token=wrong-token"):
+            pass
+    assert exc_info.value.code == 3000
+
+
 # Unit tests for terminal utility functions
 
 
@@ -404,7 +422,7 @@ def test_should_close_on_command_case_variations(self) -> None:
 @pytest.mark.skipif(is_windows, reason="Skip on Windows")
 def test_terminal_ws_unicode_input(client: TestClient) -> None:
     """Test terminal websocket with unicode input."""
-    with client.websocket_connect("/terminal/ws") as websocket:
+    with client.websocket_connect(TERMINAL_WS_URL) as websocket:
         # Send unicode command
         websocket.send_text("echo 'Hello 🌍'")
         websocket.send_text("\r")
@@ -423,7 +441,7 @@ def test_terminal_ws_invalid_session_mode(client: TestClient) -> None:
         # Test with RUN mode
         session_manager.mode = SessionMode.RUN
         with pytest.raises(WebSocketDisconnect):
-            with client.websocket_connect("/terminal/ws"):
+            with client.websocket_connect(TERMINAL_WS_URL):
                 pass  # Should fail immediately
 
     finally: