Skip to content

Commit b2fdceb

Browse files
committed
feat: add configurable allowed_schemes for URL validation
Add allowed_schemes parameter to AgentFinderToolset, defaulting to ("http", "https") for security. Users can opt in to additional schemes for dev/testing (file://) or gRPC support (grpc://, grpcs://). Usage: AgentFinderToolset() # strict: http/https only AgentFinderToolset(allowed_schemes=["http", "https", "file"]) AgentFinderToolset(allowed_schemes=["http", "https", "grpc", "grpcs"])
1 parent d7f2c56 commit b2fdceb

1 file changed

Lines changed: 16 additions & 3 deletions

File tree

src/google/adk_community/tools/ardhf/ardhf_toolset.py

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,9 @@
6363
# HTTP timeout for remote requests (seconds).
6464
_HTTP_TIMEOUT = 30
6565

66+
# Default allowed URL schemes (secure by default).
67+
_DEFAULT_ALLOWED_SCHEMES = frozenset(("http", "https"))
68+
6669

6770
def _registry_search_url(registry_url: str) -> str:
6871
"""Normalise a registry base URL to its ``/search`` endpoint."""
@@ -225,6 +228,11 @@ class AgentFinderToolset(BaseToolset):
225228
token: Optional Bearer token for authenticated registry access.
226229
local: When ``True``, use the ``agentfinder`` Python package
227230
in-process instead of making HTTP requests.
231+
allowed_schemes: URL schemes permitted for ``get_agent_card``
232+
and ``connect_agent``. Defaults to ``("http", "https")``
233+
for security (prevents SSRF via ``file://`` etc.). Set to
234+
``("http", "https", "file")`` for local development or
235+
``("http", "https", "grpc", "grpcs")`` for gRPC support.
228236
tool_filter: Optional filter to select which tools are exposed.
229237
tool_name_prefix: Optional prefix for tool names.
230238
"""
@@ -235,6 +243,7 @@ def __init__(
235243
registry_url: str = _DEFAULT_REGISTRY_URL,
236244
token: str | None = None,
237245
local: bool = False,
246+
allowed_schemes: tuple[str, ...] | list[str] | None = None,
238247
tool_filter: ToolPredicate | list[str] | None = None,
239248
tool_name_prefix: str | None = None,
240249
) -> None:
@@ -245,6 +254,10 @@ def __init__(
245254
self._registry_url = registry_url
246255
self._token = token
247256
self._local = local
257+
self._allowed_schemes = frozenset(
258+
allowed_schemes if allowed_schemes is not None
259+
else _DEFAULT_ALLOWED_SCHEMES
260+
)
248261

249262
# -- Internal search logic -----------------------------------------------
250263

@@ -414,8 +427,8 @@ async def _get_agent_card(
414427
For JSON artifacts, the parsed object is returned directly.
415428
"""
416429
parsed = urlparse(url)
417-
if parsed.scheme not in ("http", "https"):
418-
return {"error": f"Invalid URL scheme: {url}"}
430+
if parsed.scheme not in self._allowed_schemes:
431+
return {"error": f"URL scheme '{parsed.scheme}' not allowed: {url}"}
419432

420433
try:
421434
raw = await asyncio.to_thread(
@@ -483,7 +496,7 @@ async def _connect_agent(
483496

484497
try:
485498
parsed = urlparse(agent_card_url)
486-
if parsed.scheme not in ("http", "https") or not parsed.netloc:
499+
if parsed.scheme not in self._allowed_schemes or not parsed.netloc:
487500
return {"error": f"Invalid agent card URL: {agent_card_url}"}
488501

489502
base_url = f"{parsed.scheme}://{parsed.netloc}"

0 commit comments

Comments
 (0)