fix: verify datasource access and bound external observability connections

agent/Dockerfile

@@ -18,11 +18,13 @@ RUN pip install --no-cache-dir \
     # CLICKHOUSE_* env contract, or the exposed tool surface with no review signal.
     # Keep in sync with agent/requirements.txt.
     "mcp-clickhouse==0.4.1" \
-    httpx \
+    "httpx==0.28.1" \
+    "httpcore==1.0.9" \
     httpx-sse \
     aws-requests-auth
 
 COPY streamable_http_sigv4.py .
+COPY pinned_http.py .
 COPY account_utils.py .
 COPY agent.py .
 COPY readiness.py .

agent/agent.py

@@ -515,6 +515,7 @@ def _ip_always_blocked(ip_str):
     i.e. blocked REGARDLESS of the per-account allowPrivate opt-in."""
     try:
         ip = ipaddress.ip_address(ip_str)
+        ip = getattr(ip, "ipv4_mapped", None) or getattr(ip, "sixtofour", None) or ip
         if ip in _METADATA_IPS:
             return True
         return (ip.is_loopback or ip.is_link_local or ip.is_multicast or
@@ -527,6 +528,7 @@ def _ip_is_private(ip_str):
     """Returns True if the IP is in a private range (RFC1918/ULA) AND not always blocked."""
     try:
         ip = ipaddress.ip_address(ip_str)
+        ip = getattr(ip, "ipv4_mapped", None) or getattr(ip, "sixtofour", None) or ip
         # ipaddress.is_private includes link-local/loopback in some versions; we only want
         # the "opt-in" ranges (RFC1918 + RFC4193 ULA).
         return ip.is_private and not _ip_always_blocked(ip_str)
@@ -538,11 +540,10 @@ def _assert_host_allowed(url, allow_private, resolver=socket.getaddrinfo):
     """Assert the URL host is connectable: require https, resolve the host, and block any resolved IP
     that is always-blocked (metadata/loopback/link-local/...) or private-without-the-opt-in.
 
-    KNOWN LIMITATION (P4 gate, accepted — matches ADR-011 / web/lib/ssrf-guard.ts, both resolve-based):
-    this is resolve-and-recheck; the transport re-resolves at connect time, so a DNS-rebinding host
-    could return a safe IP here and a blocked IP at connect. The layered mitigations are redirect:'manual'
-    (set on both integration transports) + the registration-time literal-host guard (ssrf-guard.ts). True
-    IP-pinning (connect to the validated IP with SNI/Host preserved) is deferred to P3 hardening."""
+    Return all validated IPs for the integration transport, which pins connections while
+    preserving Host, TLS SNI and certificate verification."""
+    if "\\" in str(url):
+        raise SsrfBlocked("SSRF block: endpoint must not contain a backslash")
     parsed = urllib.parse.urlparse(url)
     if parsed.scheme != 'https':
         raise SsrfBlocked(f"SSRF block: HTTPS required for integration endpoints, got {parsed.scheme}")
@@ -553,19 +554,23 @@ def _assert_host_allowed(url, allow_private, resolver=socket.getaddrinfo):
 
     try:
         # Resolve all IPs for the host (could be dual-stack or have multiple A records)
-        addr_info = resolver(host, 443, proto=socket.IPPROTO_TCP)
+        addr_info = resolver(host, parsed.port or 443, proto=socket.IPPROTO_TCP)
     except socket.gaierror as e:
         raise SsrfBlocked(f"SSRF block: could not resolve host {host}: {e}")
 
     if not addr_info:
         raise SsrfBlocked(f"SSRF block: could not resolve host {host} (no addresses returned)")
 
+    addresses = []
     for family, type, proto, canonname, sockaddr in addr_info:
         ip_str = sockaddr[0]
         if _ip_always_blocked(ip_str):
             raise SsrfBlocked(f"SSRF block: host {host} resolved to always-blocked IP {ip_str}")
         if _ip_is_private(ip_str) and not allow_private:
             raise SsrfBlocked(f"SSRF block: host {host} resolved to private IP {ip_str} and private access disabled")
+        if ip_str not in addresses:
+            addresses.append(ip_str)
+    return tuple(addresses)
 
 
 def parse_secret(val):
@@ -589,12 +594,18 @@ def auth_headers(transport, secret):
         val = secret.get("value", secret.get("api_key", secret.get("_raw")))
         if not val:
             raise ValueError(f"Integration [api_key] missing value in secret")
+        if (not isinstance(header, str) or not re.fullmatch(r"[A-Za-z0-9!#$%&'*+.^_`|~-]+", header)
+                or header.lower() in ("host", "content-length", "transfer-encoding", "connection")
+                or not isinstance(val, str) or any(ord(ch) < 32 or ord(ch) > 126 for ch in val)):
+            raise ValueError("Integration [api_key] header configuration is invalid")
         return {header: val}
     if transport == 'oauth_client_credentials':
         # Bearer-style for this increment: {"token": "..."} or raw
         val = secret.get("token", secret.get("_raw"))
         if not val:
             raise ValueError(f"Integration [oauth_client_credentials] missing token in secret")
+        if not isinstance(val, str) or any(ord(ch) < 32 or ord(ch) > 126 for ch in val):
+            raise ValueError("Integration stored bearer token is invalid")
         return {"Authorization": f"Bearer {val}"}
     if transport == 'sigv4':
         return {}
@@ -695,21 +706,23 @@ def _connect_integration(spec, stack):
     transport = spec.get("transport")
     allow_private = bool(spec.get("allowPrivate"))
     exposed = spec.get("exposedTools") or []
-    # Connection-time SSRF (https + DNS resolve-and-recheck + private opt-in) BEFORE any network call.
-    _assert_host_allowed(endpoint, allow_private)
+    # The validated address set is request-local; warm invocations cannot share DNS pins.
+    pinned_ips = _assert_host_allowed(endpoint, allow_private)
     if transport == 'sigv4':
         # service is REQUIRED (sigv4_params raises if absent) — never silently reuse the gateway signer.
         service, region = sigv4_params(endpoint, spec.get("sigv4Service"), spec.get("sigv4Region"))
         ak, sk, tok = get_aws_credentials()
         creds = Credentials(access_key=ak, secret_key=sk, token=tok)
         client = MCPClient(lambda: streamablehttp_client_with_sigv4(
             url=endpoint, credentials=creds, service=service, region=region,
-            timeout=INTEGRATION_CONNECT_TIMEOUT))
+            timeout=INTEGRATION_CONNECT_TIMEOUT, sse_read_timeout=INTEGRATION_CONNECT_TIMEOUT,
+            pinned_ips=pinned_ips))
     else:
         secret = parse_secret(_get_secret(spec["credentialsRef"])) if spec.get("credentialsRef") else {}
         headers = auth_headers(transport, secret)  # ValueError on unknown transport
         client = MCPClient(lambda: streamablehttp_client_with_headers(
-            url=endpoint, headers=headers, timeout=INTEGRATION_CONNECT_TIMEOUT))
+            url=endpoint, headers=headers, timeout=INTEGRATION_CONNECT_TIMEOUT,
+            sse_read_timeout=INTEGRATION_CONNECT_TIMEOUT, pinned_ips=pinned_ips))
     stack.enter_context(client)
     live = get_all_tools(client)
     selected = select_integration_tools(live, exposed)

agent/lambda/clickhouse_mcp.py

@@ -273,8 +273,15 @@ def clickhouse_schema(args):
 
 
 def clickhouse_health(args):
-    """Connectivity probe for the pre-save Test / status badge: GET /ping."""
-    return ok(health(load_datasource(SLUG), "/ping"))
+    """Verify read-only SQL access using one constant row and a one-second ceiling."""
+    try:
+        result = _run_sql("SELECT 1 AS probe", 1, max_execution_time=1)
+        payload = json.loads(result["body"])
+        verified = result["statusCode"] == 200 and payload.get("collectionStatus") in ("ok", "partial")
+        return ok({"ok": verified, "verified": ["query"] if verified else [],
+                   **({} if verified else {"error": "ClickHouse query access could not be verified"})})
+    except Exception:
+        return ok({"ok": False, "error": "ClickHouse query access could not be verified"})
 
 
 _TOOLS = {

agent/lambda/datadog_mcp.py

@@ -16,7 +16,7 @@
 from cross_account import resolve_tool_name
 from datasource_http import (
     NotConnected, SsrfBlocked, assert_host_allowed, auth_headers, http_json, load_datasource,
-    set_request_conn,
+    set_request_conn, request_timeout, finite_number, observed_payload,
 )
 
 SLUG = "datadog"
@@ -45,9 +45,9 @@ def _headers(creds):
     """DD dual-key headers: dedicated apiKey/appKey fields win; else the generic custom-header pair
     (headerName/headerValue [+2]) from auth_headers covers users who configured them directly."""
     h = dict(auth_headers(creds))
-    if creds.get("apiKey"):
+    if creds.get("authType") in (None, "custom_header") and creds.get("apiKey"):
         h["DD-API-KEY"] = str(creds["apiKey"])
-    if creds.get("appKey"):
+    if creds.get("authType") in (None, "custom_header") and creds.get("appKey"):
         h["DD-APPLICATION-KEY"] = str(creds["appKey"])
     return h
 
@@ -60,11 +60,11 @@ def _ds():
 
 def _get(creds, path, params=None):
     url = creds["endpoint"].rstrip("/") + path + ("?" + urlencode(params, doseq=True) if params else "")
-    status, data = http_json("GET", url, headers=_headers(creds))
+    status, data = http_json("GET", url, headers=_headers(creds), timeout=request_timeout(creds))
     if status >= 400:
         errs = data.get("errors") if isinstance(data, dict) else None
         raise _ApiError(f"Datadog HTTP {status}: {str(errs or (data.get('raw') if isinstance(data, dict) else data))[:300]}")
-    return data
+    return observed_payload(data, status)
 
 
 def datadog_query(args):
@@ -77,18 +77,38 @@ def datadog_query(args):
         "to": _parse_time_s(args.get("end")),
     }
     data = _get(_ds(), "/api/v1/query", params)
-    series = data.get("series", []) if isinstance(data, dict) else []
+    if isinstance(data, dict) and (data.get("status") == "error" or data.get("errors") or data.get("error")):
+        return ok({"series": [], "truncated": False, "collectionStatus": "error"})
+    series = data.get("series") if isinstance(data, dict) else None
+    if not isinstance(data, dict) or data.get("status") != "ok" or not isinstance(series, list):
+        return ok({"series": [], "truncated": False, "collectionStatus": "unknown"})
     out = []
+    truncated = len(series) > MAX_SERIES
+    invalid = False
     for s in series[:MAX_SERIES]:
-        if not isinstance(s, dict):
+        if not isinstance(s, dict) or not isinstance(s.get("metric"), str) or not isinstance(s.get("pointlist"), list):
+            invalid = True
             continue
-        pts = s.get("pointlist") or []
+        pts = s["pointlist"]
+        truncated = truncated or len(pts) > MAX_POINTS
+        points = []
+        for point in pts[:MAX_POINTS]:
+            if (not isinstance(point, list) or len(point) != 2
+                    or not finite_number(point[0], 8640000000000000)
+                    or point[1] is not None and not finite_number(point[1])):
+                invalid = True
+                continue
+            points.append(point)
+        units = s.get("unit")
+        unit = units[0].get("name") if isinstance(units, list) and units and isinstance(units[0], dict) else None
         out.append({
             "metric": s.get("metric", ""), "scope": s.get("scope", ""),
-            "unit": (s.get("unit") or [{}])[0].get("name") if isinstance(s.get("unit"), list) and s.get("unit") else None,
-            "pointlist": pts[:MAX_POINTS],
+            "unit": unit, "pointlist": points,
         })
-    return ok({"series": out, "truncated": len(series) > MAX_SERIES,
+    http_status = getattr(data, "http_status", 200)
+    state = "unknown" if invalid or http_status not in (200, 206) else "partial" if truncated or http_status == 206 else "ok" if out else "empty"
+    return ok({"series": out, "truncated": truncated,
+               "collectionStatus": state,
                "status": data.get("status") if isinstance(data, dict) else None})
 
 
@@ -101,14 +121,27 @@ def datadog_schema(args):
 
 
 def datadog_health(args):
-    """Key validation probe: GET /api/v1/validate (API key), then a 1-metric list (app key scope)."""
+    """Validate configured API keys, then metric-query access (key pair or bearer)."""
     creds = load_datasource(SLUG)
     assert_host_allowed(creds["endpoint"])
     try:
-        v = _get(creds, "/api/v1/validate")
-        return ok({"ok": bool(isinstance(v, dict) and v.get("valid"))})
-    except _ApiError as e:
-        return ok({"ok": False, "error": str(e)})
+        # The API-key validator does not validate bearer access tokens. Those use
+        # the supported metric-query operation itself as the authentication proof.
+        if any(name.lower() == "dd-api-key" for name in _headers(creds)):
+            v = _get(creds, "/api/v1/validate")
+            if not isinstance(v, dict) or v.get("valid") is not True:
+                return ok({"ok": False, "error": "Datadog API key validation failed"})
+        end = int(time.time())
+        # A one-minute, single-metric query with a probe-only host filter verifies
+        # metric-query permission without fetching the fleet's real metric series.
+        probe = _get(creds, "/api/v1/query", {
+            "query": "avg:system.cpu.user{host:awsops-connection-probe}",
+            "from": str(end - 60), "to": str(end),
+        })
+        verified = isinstance(probe, dict) and probe.get("status") == "ok" and isinstance(probe.get("series"), list)
+        return ok({"ok": verified, **({} if verified else {"error": "Datadog metric-query access could not be verified"})})
+    except _ApiError:
+        return ok({"ok": False, "error": "Datadog credentials or metric-query access could not be verified"})
 
 
 _TOOLS = {

agent/lambda/datasource_http.py

@@ -24,9 +24,8 @@
 calls assert_host_allowed(endpoint) immediately before http_json(same endpoint) — that invariant is
 what makes the cache safe; it is reset every invocation by set_request_conn() (called at the top of
 every lambda_handler) so a warm container never reuses a stale pinned IP across invocations.
-(A structurally identical resolve-and-recheck gap exists in agent/agent.py's `_assert_host_allowed`
-for the separate egress-integrations transport — untouched here; that's a different call path not
-exercised by the pentest's datasource/schema-fetch findings, and needs its own fix.)
+The separate custom-MCP transport pins its request-local validated address set through
+agent/pinned_http.py; the two transports retain their different private-network contracts.
 """
 import base64
 import functools
@@ -36,11 +35,14 @@
 import os
 import re
 import socket
+import time
 import urllib.error
 import urllib.request
 from urllib.parse import urlparse
 
 HTTP_TIMEOUT = 12
+MAX_RESPONSE_BYTES = 4 * 1024 * 1024
+MAX_ERROR_BYTES = 64 * 1024
 
 _SM = None
 _SECRET_CACHE = None
@@ -62,6 +64,46 @@ class SsrfBlocked(Exception):
     """Raised when an endpoint host/scheme is disallowed."""
 
 
+class ResponseTooLarge(ValueError):
+    """The upstream body exceeded the read budget; this is never an empty result."""
+
+
+def _seconds(value, default, ceiling=55):
+    if isinstance(value, bool):
+        return default
+    text = str(value).strip()
+    if not re.fullmatch(r"\d+s?", text):
+        return default
+    return max(1, min(ceiling, int(text.removesuffix("s"))))
+
+
+def request_timeout(creds):
+    """Saved transport deadline for APIs without a provider execution-time parameter."""
+    return _seconds(creds.get("timeoutS"), HTTP_TIMEOUT)
+
+
+def execution_timeout(creds, requested=None):
+    """Provider execution ceiling, with a caller allowed only to tighten it."""
+    configured = _seconds(creds.get("timeoutS"), 10)
+    return min(configured, _seconds(requested, configured))
+
+
+def finite_number(value, bound=1.7976931348623157e308):
+    """JSON chart numbers exclude booleans, non-finite values and oversized integers."""
+    return type(value) in (int, float) and abs(value) <= bound
+
+
+class ObservedPayload(dict):
+    """Transport evidence is an attribute, so upstream JSON cannot impersonate it."""
+    def __init__(self, data, status):
+        super().__init__(data)
+        self.http_status = status
+
+
+def observed_payload(data, status):
+    return ObservedPayload(data, status) if isinstance(data, dict) else data
+
+
 # Cloud instance-metadata endpoints blocked even though private is otherwise allowed (mirror agent.py).
 _METADATA_IPS = frozenset({
     ipaddress.ip_address("169.254.169.254"),
@@ -230,15 +272,18 @@ def load_datasource(slug, instance_id=None):
        1. the request-scoped inline conn-config (`_REQUEST_CONN`) — the trusted BFF path only;
        2. the per-instance secret key `map[str(instance_id)]` — the credential-blind worker path,
           where the caller sends ONLY an instance_id and the connector reads the secret server-side;
-       3. the `map[slug]` kind-mirror (the default instance) as the fallback.
+       3. the `map[slug]` kind-mirror, only when no explicit instance was requested.
     `instance_id` is a pure argument (no module-global), so warm-container reuse cannot bleed one
     invocation's instance into the next."""
     if _REQUEST_CONN is not None:
         return _REQUEST_CONN
     m = _load_secret_map()
     creds = m.get(str(instance_id)) if instance_id is not None else None
-    if not (isinstance(creds, dict) and creds.get("endpoint")):
-        creds = m.get(slug)  # kind-mirror (default instance) fallback
+    if instance_id is not None:
+        if not isinstance(creds, dict) or not creds.get("endpoint"):
+            raise NotConnected("requested datasource instance is not configured")
+    else:
+        creds = m.get(slug)
     if not isinstance(creds, dict) or not creds.get("endpoint"):
         raise NotConnected(f"{slug} not connected (no endpoint configured in the Connectors UI)")
     return creds
@@ -341,18 +386,36 @@ def _parse(raw):
         return {"raw": raw.decode("utf-8", "replace")[:4000]}
 
 
-def http_json(method, url, headers=None, body=None, timeout=HTTP_TIMEOUT):
+def _read_bounded(response, max_bytes, timeout):
+    """Read at most budget+1 bytes; never trust Content-Length or parse before the bound."""
+    chunks = []
+    total = 0
+    deadline = time.monotonic() + timeout
+    read = getattr(response, "read1", response.read)
+    try:
+        while True:
+            if time.monotonic() >= deadline:
+                raise TimeoutError("datasource response read deadline exceeded")
+            chunk = read(min(65536, max_bytes - total + 1))
+            if not chunk:
+                return b"".join(chunks)
+            total += len(chunk)
+            if total > max_bytes:
+                raise ResponseTooLarge(f"datasource response exceeds {max_bytes} bytes")
+            chunks.append(chunk)
+    finally:
+        response.close()
+
+
+def http_json(method, url, headers=None, body=None, timeout=HTTP_TIMEOUT, max_bytes=MAX_RESPONSE_BYTES):
     """Send a request (no auto-redirect). Returns (status, parsed_dict). Non-2xx → (status, body).
     Pinned to the IPs assert_host_allowed(url) most recently validated for this host, if any (tried
     in order with fallback) — see module docstring (IP-pinning)."""
     data = body if isinstance(body, (bytes, bytearray)) else (body.encode() if isinstance(body, str) else None)
     req = urllib.request.Request(url, data=data, method=method, headers=headers or {})
     try:
         resp = _opener_for(urlparse(url).hostname).open(req, timeout=timeout)
-        return resp.getcode(), _parse(resp.read())
+        return resp.getcode(), _parse(_read_bounded(resp, max_bytes, timeout))
     except urllib.error.HTTPError as e:
-        try:
-            raw = e.read()
-        except Exception:  # noqa: BLE001
-            raw = b""
+        raw = _read_bounded(e, min(max_bytes, MAX_ERROR_BYTES), timeout)
         return e.code, _parse(raw)