fix: add protection for arbitrary module imports

Close#4947
Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 888296476

Author: GWeale

src/google/adk/cli/adk_web_server.py

@@ -20,6 +20,7 @@
 importjson
 importlogging
 importos
+importre
 importsys
 importtime
 importtraceback
@@ -138,6 +139,158 @@ def _parse_cors_origins(
 returnliteral_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"*"inallowed_literal_origins:
+returnTrue
+iforigininallowed_literal_origins:
+returnTrue
+ifallowed_origin_regexisnotNone:
+returnallowed_origin_regex.fullmatch(origin) isnotNone
+returnFalse
+
+
+def_normalize_origin_scheme(scheme: str) ->str:
+"""Normalize request schemes to the browser Origin scheme space."""
+ifscheme=="ws":
+return"http"
+ifscheme=="wss":
+return"https"
+returnscheme
+
+
+def_strip_optional_quotes(value: str) ->str:
+"""Strip a single pair of wrapping quotes from a header value."""
+iflen(value) >=2andvalue[0] =='"'andvalue[-1] =='"':
+returnvalue[1:-1]
+returnvalue
+
+
+def_get_scope_header(
+scope: dict[str, Any], header_name: bytes
+) ->Optional[str]:
+"""Return the first matching header value from an ASGI scope."""
+forcandidate_name, candidate_valueinscope.get("headers", []):
+ifcandidate_name==header_name:
+returncandidate_value.decode("latin-1").split(",", 1)[0].strip()
+returnNone
+
+
+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")
+ifforwardedisnotNone:
+proto=None
+host=None
+forelementinforwarded.split(",", 1)[0].split(";"):
+if"="notinelement:
+continue
+name, value=element.split("=", 1)
+ifname.strip().lower() =="proto":
+proto=_strip_optional_quotes(value.strip())
+elifname.strip().lower() =="host":
+host=_strip_optional_quotes(value.strip())
+ifprotoisnotNoneandhostisnotNone:
+returnf"{_normalize_origin_scheme(proto)}://{host}"
+
+host=_get_scope_header(scope, b"x-forwarded-host")
+ifhostisNone:
+host=_get_scope_header(scope, b"host")
+ifhostisNone:
+returnNone
+
+proto=_get_scope_header(scope, b"x-forwarded-proto")
+ifprotoisNone:
+proto=scope.get("scheme", "http")
+returnf"{_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."""
+ifhas_configured_allowed_originsand_is_origin_allowed(
+origin, allowed_literal_origins, allowed_origin_regex
+ ):
+returnTrue
+
+request_origin=_get_request_origin(scope)
+ifrequest_originisNone:
+returnFalse
+returnorigin==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
+
+asyncdef__call__(
+self,
+scope: dict[str, Any],
+receive: Any,
+send: Any,
+ ) ->None:
+ifscope["type"] !="http":
+awaitself._app(scope, receive, send)
+return
+
+method=scope.get("method", "GET")
+ifmethodin_SAFE_HTTP_METHODS:
+awaitself._app(scope, receive, send)
+return
+
+origin=_get_scope_header(scope, b"origin")
+iforiginisNone:
+awaitself._app(scope, receive, send)
+return
+
+if_is_request_origin_allowed(
+origin,
+scope,
+self._allowed_origins,
+self._allowed_origin_regex,
+self._has_configured_allowed_origins,
+ ):
+awaitself._app(scope, receive, send)
+return
+
+response_body=b"Forbidden: origin not allowed"
+awaitsend({
+"type": "http.response.start",
+"status": 403,
+"headers": [
+ (b"content-type", b"text/plain"),
+ (b"content-length", str(len(response_body)).encode()),
+ ],
+ })
+awaitsend({
+"type": "http.response.body",
+"body": response_body,
+ })
+
+
 classApiServerSpanExporter(export_lib.SpanExporter):
 
 def__init__(self, trace_dict):
@@ -757,8 +910,12 @@ async def internal_lifespan(app: FastAPI):
 # Run the FastAPI server.
 app=FastAPI(lifespan=internal_lifespan)
 
+has_configured_allowed_origins=bool(allow_origins)
 ifallow_origins:
 literal_origins, combined_regex=_parse_cors_origins(allow_origins)
+compiled_origin_regex= (
+re.compile(combined_regex) ifcombined_regexisnotNoneelseNone
+ )
 app.add_middleware(
 CORSMiddleware,
 allow_origins=literal_origins,
@@ -767,6 +924,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")
 asyncdefhealth() ->dict[str, str]:
@@ -1802,14 +1969,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")
+ifws_originisnotNoneandnot_is_request_origin_allowed(
+ws_origin,
+websocket.scope,
+literal_origins,
+compiled_origin_regex,
+has_configured_allowed_origins,
+ ):
+awaitwebsocket.close(code=1008, reason="Origin not allowed")
+return
+
 awaitwebsocket.accept()
 
 session=awaitself.session_service.get_session(
 app_name=app_name, user_id=user_id, session_id=session_id
  )
 ifnotsession:
-# Accept first so that the client is aware of connection establishment,
-# then close with a specific code.
 awaitwebsocket.close(code=1002, reason="Session not found")
 return
 

tests/unittests/cli/test_adk_web_server_run_live.py

@@ -22,6 +22,7 @@
 fromgoogle.adk.events.eventimportEvent
 fromgoogle.adk.sessions.in_memory_session_serviceimportInMemorySessionService
 importpytest
+fromstarlette.websocketsimportWebSocketDisconnect
 
 
 class_DummyAgent(BaseAgent):
@@ -203,3 +204,75 @@ async def _get_runner_async(_self, _app_name: str):
 run_config.session_resumption.transparent
 isexpected_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=".",
+ )
+
+asyncdef_get_runner_async(_self, _app_name: str):
+returnrunner
+
+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,
+ )
+returnTestClient(fast_api_app)
+
+
+deftest_run_live_rejects_disallowed_origin():
+client=_build_ws_client()
+withpytest.raises(WebSocketDisconnect) asexc_info:
+withclient.websocket_connect(
+_WS_BASE_URL,
+headers={"origin": "https://evil.com"},
+ ) asws:
+ws.receive_text()
+assertexc_info.value.code==1008
+
+
+deftest_run_live_allows_matching_origin():
+client=_build_ws_client()
+withclient.websocket_connect(
+_WS_BASE_URL,
+headers={"origin": "http://testserver"},
+ ) asws:
+_=ws.receive_text()
+
+
+deftest_run_live_allows_no_origin_header():
+"""Non-browser clients (curl, wscat, SDKs) send no Origin header."""
+client=_build_ws_client()
+withclient.websocket_connect(_WS_BASE_URL) asws:
+_=ws.receive_text()

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(
 assertnottmp_dir.exists() ornotany(tmp_dir.iterdir())
 
 
+deftest_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"),
+ )],
+ )
+
+assertresponse.status_code==403
+assertresponse.text=="Forbidden: origin not allowed"
+assertnot (tmp_path/"app"/"tmp"/"app").exists()
+
+
+deftest_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"),
+ )],
+ )
+
+assertresponse.status_code==200
+assertresponse.json() isTrue
+assert (tmp_path/"app"/"tmp"/"app"/"root_agent.yaml").is_file()
+
+
+deftest_builder_get_allows_cross_origin_get(builder_test_client):
+response=builder_test_client.get(
+"/builder/app/missing?tmp=true",
+headers={"origin": "https://evil.com"},
+ )
+
+assertresponse.status_code==200
+assertresponse.text==""
+
+
 deftest_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)