diff --git a/agent_assembly/core/assembly.py b/agent_assembly/core/assembly.py index 0ffaf0b..05be85d 100644 --- a/agent_assembly/core/assembly.py +++ b/agent_assembly/core/assembly.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +import re import sys from dataclasses import dataclass, field from threading import Lock @@ -42,12 +43,27 @@ ENV_CONTROL_PLANE_URL = "AA_CONTROL_PLANE_URL" _DEFAULT_AGENT_ID = "agent-assembly-default" +_AGENT_ID_RE = re.compile(r"^[A-Za-z0-9_.-]{1,128}$") _VALID_RUNTIME_MODES = {"auto", "ebpf", "proxy", "sdk-only"} _VALID_ENFORCEMENT_MODES: frozenset[EnforcementMode] = frozenset({"enforce", "observe", "disabled"}) _INIT_LOCK = Lock() _ACTIVE_CONTEXT: AssemblyContext | None = None +def _validate_agent_id(agent_id: str) -> str: + """Validate agent_id format before it reaches socket-path interpolation (AAASM-4301). + + Only allows [A-Za-z0-9_.-] up to 128 chars. Reject anything else with a clear ValueError + so callers get a fast, actionable error instead of a downstream FileNotFoundError from + a malformed socket path. + """ + if not _AGENT_ID_RE.fullmatch(agent_id): + raise ValueError( + f"invalid agent_id: {agent_id!r}; must match {_AGENT_ID_RE.pattern}" + ) + return agent_id + + class RuntimePatch(Protocol): """Internal monkey-patch mechanism used by framework adapters. @@ -204,7 +220,7 @@ def init_assembly( if spawned_by_tool is None: spawned_by_tool = _spawn.spawned_by_tool - resolved_agent_id = agent_id or _DEFAULT_AGENT_ID + resolved_agent_id = _validate_agent_id(agent_id or _DEFAULT_AGENT_ID) global _ACTIVE_CONTEXT with _INIT_LOCK: diff --git a/test/unit/test_assembly.py b/test/unit/test_assembly.py index b087f5d..ac7e2d9 100644 --- a/test/unit/test_assembly.py +++ b/test/unit/test_assembly.py @@ -698,3 +698,95 @@ def test_init_assembly_no_warning_for_loopback_http_with_api_key( agent_id="agent-loopback", ) context.shutdown() + + +# AAASM-4301: validate_agent_id enforces a strict character-set + length cap +# on the agent_id argument before it ever reaches socket-path interpolation. + + +@pytest.mark.parametrize( + "agent_id", + [ + "agent-assembly-default", # the module default + "default", + "agent-1", + "my_agent_2026", + "a.b.c", + "A", # single char, upper + "a", # single char, lower + "0", # single char, digit + "a" * 128, # max length + ], +) +def test_validate_agent_id_accepts_well_formed_values(agent_id: str) -> None: + """Well-formed agent_ids match the regex and are returned unchanged.""" + assert core_assembly._validate_agent_id(agent_id) == agent_id + + +@pytest.mark.parametrize( + "agent_id", + [ + "", # empty string + "a" * 129, # one char over the length cap + "../etc/passwd", # path-traversal attempt + "has space", # whitespace + "has/slash", # path separator + "has$special", # shell metacharacter + "unicode-α", # non-ASCII alpha + "trailing\n", # newline + ], +) +def test_validate_agent_id_rejects_malformed_values(agent_id: str) -> None: + """Malformed agent_ids raise ValueError with the regex pattern in the message.""" + with pytest.raises(ValueError, match=r"invalid agent_id"): + core_assembly._validate_agent_id(agent_id) + + +def test_init_assembly_rejects_malformed_agent_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """AAASM-4301 regression: malformed agent_id fails fast at init_assembly entry. + + Without this validation, a path-traversal-like agent_id would be silently + interpolated into the runtime socket path, causing a downstream FileNotFoundError + when the SDK tries to connect. The validator raises ValueError up-front so the + caller gets a clear, actionable error. + """ + monkeypatch.setattr(core_assembly, "_register_adapters", lambda **kwargs: []) + monkeypatch.setattr( + core_assembly, + "_start_network_layer", + lambda **kwargs: ("sdk-only", core_assembly._noop_shutdown), + ) + + with pytest.raises(ValueError, match=r"invalid agent_id"): + init_assembly( + gateway_url="http://localhost:8080", + api_key="test-api-key", + agent_id="../attack", + ) + + +def test_init_assembly_accepts_default_agent_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The module-level _DEFAULT_AGENT_ID must satisfy the validator. + + Guards against a future change that alters _DEFAULT_AGENT_ID to a value the + regex rejects, which would break every caller of init_assembly() (no args). + """ + monkeypatch.setattr(core_assembly, "_register_adapters", lambda **kwargs: []) + monkeypatch.setattr( + core_assembly, + "_start_network_layer", + lambda **kwargs: ("sdk-only", core_assembly._noop_shutdown), + ) + + context = init_assembly( + gateway_url="http://localhost:8080", + api_key="test-api-key", + ) + try: + assert context.client.agent_id == core_assembly._DEFAULT_AGENT_ID + finally: + context.shutdown()