fix: add protection for arbitrary module imports

GWeale · sasha-gitg · commit 0019801f6080 · 2026-03-26T16:48:49.000-04:00
Close #4947 Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 888296476
diff --git a/src/google/adk/cli/adk_web_server.py b/src/google/adk/cli/adk_web_server.py
@@ -20,6 +20,7 @@
 import json
 import logging
 import os
+import re
 import sys
 import time
 import traceback
@@ -140,6 +141,158 @@ def _parse_cors_origins(
   return literal_origins, combined_regex
 
 
+def _is_origin_allowed(
+    origin: str,
+    allowed_literal_origins: list[str],
+    allowed_origin_regex: Optional[re.Pattern[str]],
+) -> bool:
+  """Check whether the given origin matches the allowed origins."""
+  if "*" in allowed_literal_origins:
+    return True
+  if origin in allowed_literal_origins:
+    return True
+  if allowed_origin_regex is not None:
+    return allowed_origin_regex.fullmatch(origin) is not None
+  return False
+
+
+def _normalize_origin_scheme(scheme: str) -> str:
+  """Normalize request schemes to the browser Origin scheme space."""
+  if scheme == "ws":
+    return "http"
+  if scheme == "wss":
+    return "https"
+  return scheme
+
+
+def _strip_optional_quotes(value: str) -> str:
+  """Strip a single pair of wrapping quotes from a header value."""
+  if len(value) >= 2 and value[0] == '"' and value[-1] == '"':
+    return value[1:-1]
+  return value
+
+
+def _get_scope_header(
+    scope: dict[str, Any], header_name: bytes
+) -> Optional[str]:
+  """Return the first matching header value from an ASGI scope."""
+  for candidate_name, candidate_value in scope.get("headers", []):
+    if candidate_name == header_name:
+      return candidate_value.decode("latin-1").split(",", 1)[0].strip()
+  return None
+
+
+def _get_request_origin(scope: dict[str, Any]) -> Optional[str]:
+  """Compute the effective origin for the current HTTP/WebSocket request."""
+  forwarded = _get_scope_header(scope, b"forwarded")
+  if forwarded is not None:
+    proto = None
+    host = None
+    for element in forwarded.split(",", 1)[0].split(";"):
+      if "=" not in element:
+        continue
+      name, value = element.split("=", 1)
+      if name.strip().lower() == "proto":
+        proto = _strip_optional_quotes(value.strip())
+      elif name.strip().lower() == "host":
+        host = _strip_optional_quotes(value.strip())
+    if proto is not None and host is not None:
+      return f"{_normalize_origin_scheme(proto)}://{host}"
+
+  host = _get_scope_header(scope, b"x-forwarded-host")
+  if host is None:
+    host = _get_scope_header(scope, b"host")
+  if host is None:
+    return None
+
+  proto = _get_scope_header(scope, b"x-forwarded-proto")
+  if proto is None:
+    proto = scope.get("scheme", "http")
+  return f"{_normalize_origin_scheme(proto)}://{host}"
+
+
+def _is_request_origin_allowed(
+    origin: str,
+    scope: dict[str, Any],
+    allowed_literal_origins: list[str],
+    allowed_origin_regex: Optional[re.Pattern[str]],
+    has_configured_allowed_origins: bool,
+) -> bool:
+  """Validate an Origin header against explicit config or same-origin."""
+  if has_configured_allowed_origins and _is_origin_allowed(
+      origin, allowed_literal_origins, allowed_origin_regex
+  ):
+    return True
+
+  request_origin = _get_request_origin(scope)
+  if request_origin is None:
+    return False
+  return origin == request_origin
+
+
+_SAFE_HTTP_METHODS = frozenset({"GET", "HEAD", "OPTIONS"})
+
+
+class _OriginCheckMiddleware:
+  """ASGI middleware that blocks cross-origin state-changing requests."""
+
+  def __init__(
+      self,
+      app: Any,
+      has_configured_allowed_origins: bool,
+      allowed_origins: list[str],
+      allowed_origin_regex: Optional[re.Pattern[str]],
+  ) -> None:
+    self._app = app
+    self._has_configured_allowed_origins = has_configured_allowed_origins
+    self._allowed_origins = allowed_origins
+    self._allowed_origin_regex = allowed_origin_regex
+
+  async def __call__(
+      self,
+      scope: dict[str, Any],
+      receive: Any,
+      send: Any,
+  ) -> None:
+    if scope["type"] != "http":
+      await self._app(scope, receive, send)
+      return
+
+    method = scope.get("method", "GET")
+    if method in _SAFE_HTTP_METHODS:
+      await self._app(scope, receive, send)
+      return
+
+    origin = _get_scope_header(scope, b"origin")
+    if origin is None:
+      await self._app(scope, receive, send)
+      return
+
+    if _is_request_origin_allowed(
+        origin,
+        scope,
+        self._allowed_origins,
+        self._allowed_origin_regex,
+        self._has_configured_allowed_origins,
+    ):
+      await self._app(scope, receive, send)
+      return
+
+    response_body = b"Forbidden: origin not allowed"
+    await send({
+        "type": "http.response.start",
+        "status": 403,
+        "headers": [
+            (b"content-type", b"text/plain"),
+            (b"content-length", str(len(response_body)).encode()),
+        ],
+    })
+    await send({
+        "type": "http.response.body",
+        "body": response_body,
+    })
+
+
 class ApiServerSpanExporter(export_lib.SpanExporter):
 
   def __init__(self, trace_dict):
@@ -759,8 +912,12 @@ async def internal_lifespan(app: FastAPI):
     # Run the FastAPI server.
     app = FastAPI(lifespan=internal_lifespan)
 
+    has_configured_allowed_origins = bool(allow_origins)
     if allow_origins:
       literal_origins, combined_regex = _parse_cors_origins(allow_origins)
+      compiled_origin_regex = (
+          re.compile(combined_regex) if combined_regex is not None else None
+      )
       app.add_middleware(
           CORSMiddleware,
           allow_origins=literal_origins,
@@ -769,6 +926,16 @@ async def internal_lifespan(app: FastAPI):
           allow_methods=["*"],
           allow_headers=["*"],
       )
+    else:
+      literal_origins = []
+      compiled_origin_regex = None
+
+    app.add_middleware(
+        _OriginCheckMiddleware,
+        has_configured_allowed_origins=has_configured_allowed_origins,
+        allowed_origins=literal_origins,
+        allowed_origin_regex=compiled_origin_regex,
+    )
 
     @app.get("/health")
     async def health() -> dict[str, str]:
@@ -1755,14 +1922,23 @@ async def run_agent_live(
         enable_affective_dialog: bool | None = Query(default=None),
         enable_session_resumption: bool | None = Query(default=None),
     ) -> None:
+      ws_origin = websocket.headers.get("origin")
+      if ws_origin is not None and not _is_request_origin_allowed(
+          ws_origin,
+          websocket.scope,
+          literal_origins,
+          compiled_origin_regex,
+          has_configured_allowed_origins,
+      ):
+        await websocket.close(code=1008, reason="Origin not allowed")
+        return
+
       await websocket.accept()
 
       session = await self.session_service.get_session(
           app_name=app_name, user_id=user_id, session_id=session_id
       )
       if not session:
-        # Accept first so that the client is aware of connection establishment,
-        # then close with a specific code.
         await websocket.close(code=1002, reason="Session not found")
         return
 
diff --git a/tests/unittests/cli/test_adk_web_server_run_live.py b/tests/unittests/cli/test_adk_web_server_run_live.py
@@ -22,6 +22,7 @@
 from google.adk.events.event import Event
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
 import pytest
+from starlette.websockets import WebSocketDisconnect
 
 
 class _DummyAgent(BaseAgent):
@@ -203,3 +204,75 @@ async def _get_runner_async(_self, _app_name: str):
         run_config.session_resumption.transparent
         is expected_session_resumption_transparent
     )
+
+
+_WS_BASE_URL = (
+    "/run_live"
+    "?app_name=test_app"
+    "&user_id=user"
+    "&session_id=session"
+    "&modalities=AUDIO"
+)
+
+
+def _build_ws_client():
+  """Build a TestClient wired to a capturing runner."""
+  session_service = InMemorySessionService()
+  asyncio.run(
+      session_service.create_session(
+          app_name="test_app",
+          user_id="user",
+          session_id="session",
+          state={},
+      )
+  )
+
+  runner = _CapturingRunner()
+  adk_web_server = AdkWebServer(
+      agent_loader=_DummyAgentLoader(),
+      session_service=session_service,
+      memory_service=types.SimpleNamespace(),
+      artifact_service=types.SimpleNamespace(),
+      credential_service=types.SimpleNamespace(),
+      eval_sets_manager=types.SimpleNamespace(),
+      eval_set_results_manager=types.SimpleNamespace(),
+      agents_dir=".",
+  )
+
+  async def _get_runner_async(_self, _app_name: str):
+    return runner
+
+  adk_web_server.get_runner_async = _get_runner_async.__get__(adk_web_server)  # pytype: disable=attribute-error
+
+  fast_api_app = adk_web_server.get_fast_api_app(
+      setup_observer=lambda _observer, _server: None,
+      tear_down_observer=lambda _observer, _server: None,
+  )
+  return TestClient(fast_api_app)
+
+
+def test_run_live_rejects_disallowed_origin():
+  client = _build_ws_client()
+  with pytest.raises(WebSocketDisconnect) as exc_info:
+    with client.websocket_connect(
+        _WS_BASE_URL,
+        headers={"origin": "https://evil.com"},
+    ) as ws:
+      ws.receive_text()
+  assert exc_info.value.code == 1008
+
+
+def test_run_live_allows_matching_origin():
+  client = _build_ws_client()
+  with client.websocket_connect(
+      _WS_BASE_URL,
+      headers={"origin": "http://testserver"},
+  ) as ws:
+    _ = ws.receive_text()
+
+
+def test_run_live_allows_no_origin_header():
+  """Non-browser clients (curl, wscat, SDKs) send no Origin header."""
+  client = _build_ws_client()
+  with client.websocket_connect(_WS_BASE_URL) as ws:
+    _ = ws.receive_text()
diff --git a/tests/unittests/cli/test_fast_api.py b/tests/unittests/cli/test_fast_api.py
@@ -593,7 +593,7 @@ def builder_test_client(
         session_service_uri="",
         artifact_service_uri="",
         memory_service_uri="",
-        allow_origins=["*"],
+        allow_origins=None,
         a2a=False,
         host="127.0.0.1",
         port=8000,
@@ -1595,6 +1595,46 @@ def test_builder_final_save_preserves_tools_and_cleans_tmp(
   assert not tmp_dir.exists() or not any(tmp_dir.iterdir())
 
 
+def test_builder_save_rejects_cross_origin_post(builder_test_client, tmp_path):
+  response = builder_test_client.post(
+      "/builder/save?tmp=true",
+      headers={"origin": "https://evil.com"},
+      files=[(
+          "files",
+          ("app/root_agent.yaml", b"name: app\n", "application/x-yaml"),
+      )],
+  )
+
+  assert response.status_code == 403
+  assert response.text == "Forbidden: origin not allowed"
+  assert not (tmp_path / "app" / "tmp" / "app").exists()
+
+
+def test_builder_save_allows_same_origin_post(builder_test_client, tmp_path):
+  response = builder_test_client.post(
+      "/builder/save?tmp=true",
+      headers={"origin": "http://testserver"},
+      files=[(
+          "files",
+          ("app/root_agent.yaml", b"name: app\n", "application/x-yaml"),
+      )],
+  )
+
+  assert response.status_code == 200
+  assert response.json() is True
+  assert (tmp_path / "app" / "tmp" / "app" / "root_agent.yaml").is_file()
+
+
+def test_builder_get_allows_cross_origin_get(builder_test_client):
+  response = builder_test_client.get(
+      "/builder/app/missing?tmp=true",
+      headers={"origin": "https://evil.com"},
+  )
+
+  assert response.status_code == 200
+  assert response.text == ""
+
+
 def test_builder_cancel_deletes_tmp_idempotent(builder_test_client, tmp_path):
   tmp_agent_root = tmp_path / "app" / "tmp" / "app"
   tmp_agent_root.mkdir(parents=True, exist_ok=True)
