zeph_a2a/
client.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! A2A protocol HTTP client with optional TLS enforcement and SSRF protection.
5
6use std::net::SocketAddr;
7use std::pin::Pin;
8use std::time::Duration;
9
10use eventsource_stream::Eventsource;
11use futures_core::Stream;
12use serde::{Deserialize, Serialize, de::DeserializeOwned};
13use tokio_stream::StreamExt;
14use zeph_common::net::resolve_and_validate;
15
16use crate::error::A2aError;
17use crate::ibct::{Ibct, IbctKey, ibct_scope_origin};
18use crate::jsonrpc::{
19    JsonRpcRequest, JsonRpcResponse, METHOD_CANCEL_TASK, METHOD_GET_TASK, METHOD_SEND_MESSAGE,
20    METHOD_SEND_STREAMING_MESSAGE, SendMessageParams, TaskIdParams,
21};
22use crate::types::{Task, TaskArtifactUpdateEvent, TaskStatusUpdateEvent};
23
24/// A pinned, heap-allocated stream of [`TaskEvent`]s from a streaming A2A call.
25///
26/// Produced by [`A2aClient::stream_message`]. Each item is either a status update
27/// or an artifact update; errors are surfaced inline as `Err(A2aError)`.
28pub type TaskEventStream = Pin<Box<dyn Stream<Item = Result<TaskEvent, A2aError>> + Send>>;
29
30/// A single event received on a streaming (`message/stream`) A2A connection.
31///
32/// The A2A spec multiplexes two event kinds over the same SSE channel. This enum
33/// uses `#[serde(untagged)]` so that the deserializer inspects the `kind` field
34/// inside the inner struct to determine the variant.
35#[non_exhaustive]
36#[derive(Debug, Clone, Serialize, Deserialize)]
37#[serde(untagged)]
38pub enum TaskEvent {
39    /// A task lifecycle transition (e.g., `submitted` → `working` → `completed`).
40    StatusUpdate(TaskStatusUpdateEvent),
41    /// A new or updated output artifact from the agent.
42    ArtifactUpdate(TaskArtifactUpdateEvent),
43}
44
45/// Security posture applied to outbound [`A2aClient`] requests.
46///
47/// Named fields eliminate the transposition hazard of a two-bool builder method
48/// (`with_security(true, false)` vs. `with_security(false, true)` are easy to swap
49/// by accident) and group the security boundary as one reviewable unit.
50///
51/// # Examples
52///
53/// ```rust
54/// use zeph_a2a::{A2aClient, SecurityPolicy};
55///
56/// // Recommended for production: reject HTTP and private/loopback targets.
57/// let client = A2aClient::new(reqwest::Client::new()).with_security(SecurityPolicy::hardened());
58///
59/// // Partial policy via named fields — no ambiguity about which flag is which.
60/// let tls_only = SecurityPolicy {
61///     require_tls: true,
62///     ssrf_protection: false,
63/// };
64/// let _ = client.with_security(tls_only);
65/// ```
66#[derive(Debug, Clone, Copy, PartialEq, Eq)]
67pub struct SecurityPolicy {
68    /// Reject any endpoint that does not start with `https://`, and build requests with
69    /// `https_only(true)` so a redirect cannot silently downgrade the connection to `http://`.
70    pub require_tls: bool,
71    /// Resolve the endpoint hostname via DNS, reject private/loopback/link-local ranges,
72    /// and pin the validated address for the actual connection so it cannot be re-resolved
73    /// to a different (attacker-controlled) address between the check and the connect.
74    pub ssrf_protection: bool,
75}
76
77impl SecurityPolicy {
78    /// Both protections enabled. The recommended posture for production deployments
79    /// that talk to untrusted or third-party A2A endpoints.
80    ///
81    /// # Examples
82    ///
83    /// ```rust
84    /// use zeph_a2a::SecurityPolicy;
85    ///
86    /// let policy = SecurityPolicy::hardened();
87    /// assert!(policy.require_tls);
88    /// assert!(policy.ssrf_protection);
89    /// ```
90    #[must_use]
91    pub const fn hardened() -> Self {
92        Self {
93            require_tls: true,
94            ssrf_protection: true,
95        }
96    }
97
98    /// Both protections disabled. Suitable only for local development against
99    /// trusted, non-adversarial endpoints (e.g. `http://localhost`).
100    ///
101    /// # Examples
102    ///
103    /// ```rust
104    /// use zeph_a2a::SecurityPolicy;
105    ///
106    /// let policy = SecurityPolicy::permissive();
107    /// assert!(!policy.require_tls);
108    /// assert!(!policy.ssrf_protection);
109    /// ```
110    #[must_use]
111    pub const fn permissive() -> Self {
112        Self {
113            require_tls: false,
114            ssrf_protection: false,
115        }
116    }
117}
118
119/// A DNS-validated hostname and its resolved addresses, used to pin the actual HTTP
120/// connection to the exact addresses that passed SSRF validation (see [`SecurityPolicy`]).
121#[derive(Debug)]
122struct PinnedTarget {
123    host: String,
124    addrs: Vec<SocketAddr>,
125}
126
127/// HTTP client for the A2A protocol.
128///
129/// `A2aClient` wraps a `reqwest::Client` and provides typed methods for the four
130/// A2A JSON-RPC operations: `message/send`, `message/stream`, `tasks/get`, and
131/// `tasks/cancel`. Each call optionally accepts a bearer token for authentication.
132///
133/// # Security
134///
135/// Use [`with_security`](A2aClient::with_security) to harden the client for
136/// production deployments — see [`SecurityPolicy`]. When either flag is enabled,
137/// each request is sent through a dedicated per-request `reqwest::Client` with
138/// redirects disabled and, when `ssrf_protection` is on, the connection pinned to
139/// the exact addresses that were validated (no re-resolution at connect time).
140///
141/// # Examples
142///
143/// ```rust,no_run
144/// use zeph_a2a::{A2aClient, SecurityPolicy, SendMessageParams, Message};
145///
146/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
147/// let client = A2aClient::new(reqwest::Client::new())
148///     .with_security(SecurityPolicy::hardened());
149///
150/// let params = SendMessageParams {
151///     message: Message::user_text("Summarize this page."),
152///     configuration: None,
153/// };
154/// let task = client.send_message("https://agent.example.com/a2a", params, Some("tok")).await?;
155/// println!("Task state: {:?}", task.status.state);
156/// # Ok(())
157/// # }
158/// ```
159pub struct A2aClient {
160    client: reqwest::Client,
161    security: SecurityPolicy,
162    /// Per-request timeout applied to `rpc_call` (send + JSON parse) and to the initial
163    /// `send()` in `stream_message`. The SSE body stream itself is not bounded — that
164    /// is the caller's responsibility.
165    ///
166    /// If the underlying `reqwest::Client` was also built with `.timeout()`, both limits
167    /// race: whichever fires first wins. `request_timeout` takes semantic priority because
168    /// it maps to `A2aError::Timeout`; the reqwest-level timeout maps to `A2aError::Http`.
169    request_timeout: Duration,
170    /// IBCT signing key. When set, every request carries a scoped `X-Zeph-IBCT` header
171    /// alongside the bearer token (see [`with_ibct_key`](Self::with_ibct_key)).
172    ibct_key: Option<IbctKey>,
173    /// TTL for issued IBCT tokens. Has no effect unless `ibct_key` is set.
174    ibct_ttl: Duration,
175}
176
177impl A2aClient {
178    /// Create a new `A2aClient` with no security restrictions.
179    ///
180    /// Security features are disabled by default for local/dev usage. Enable them
181    /// with [`with_security`](Self::with_security) for production deployments.
182    #[must_use]
183    pub fn new(client: reqwest::Client) -> Self {
184        Self {
185            client,
186            security: SecurityPolicy::permissive(),
187            request_timeout: Duration::from_secs(30),
188            ibct_key: None,
189            ibct_ttl: Duration::from_mins(5),
190        }
191    }
192
193    /// Configure the [`SecurityPolicy`] for this client.
194    ///
195    /// Defaults to [`SecurityPolicy::permissive()`] (no restrictions). This method
196    /// uses the builder pattern and can be chained directly after [`new`](Self::new).
197    ///
198    /// # Examples
199    ///
200    /// ```rust
201    /// use zeph_a2a::{A2aClient, SecurityPolicy};
202    ///
203    /// let client = A2aClient::new(reqwest::Client::new())
204    ///     .with_security(SecurityPolicy::hardened());
205    /// ```
206    #[must_use]
207    pub fn with_security(mut self, policy: SecurityPolicy) -> Self {
208        self.security = policy;
209        self
210    }
211
212    /// Set the per-request timeout for RPC and streaming connection calls (default: 30 seconds).
213    ///
214    /// Applied to the full send + JSON response parse in `rpc_call`, and to the initial
215    /// HTTP `send()` in `stream_message`. The SSE body stream after connection is intentionally
216    /// unbounded — streams can legitimately run for a long time.
217    #[must_use]
218    pub fn with_request_timeout(mut self, timeout: Duration) -> Self {
219        self.request_timeout = timeout;
220        self
221    }
222
223    /// Configure an IBCT signing key so every request carries a scoped `X-Zeph-IBCT` header
224    /// alongside the bearer token.
225    ///
226    /// The token is scoped to the exact `endpoint` string passed to `send_message`/
227    /// `get_task`/`cancel_task`/`stream_message`, and to the request's `task_id` —
228    /// `params.id` for `get_task`/`cancel_task`, `message.task_id` for `send_message`/
229    /// `stream_message` (the empty-string sentinel when the message has no `task_id` yet,
230    /// i.e. it starts a brand-new task the server has not assigned an ID to).
231    ///
232    /// Issuance failures (e.g. this crate compiled without the `ibct` feature) are logged
233    /// and the request proceeds without the header — the server decides whether to reject it.
234    ///
235    /// # Examples
236    ///
237    /// ```rust
238    /// use zeph_a2a::{A2aClient, IbctKey};
239    ///
240    /// let key = IbctKey { key_id: "k1".into(), key_bytes: b"secret".to_vec() };
241    /// let client = A2aClient::new(reqwest::Client::new()).with_ibct_key(key);
242    /// ```
243    #[must_use]
244    pub fn with_ibct_key(mut self, key: IbctKey) -> Self {
245        self.ibct_key = Some(key);
246        self
247    }
248
249    /// Set the TTL for issued IBCT tokens (default: 5 minutes). Has no effect unless
250    /// [`with_ibct_key`](Self::with_ibct_key) is also configured.
251    #[must_use]
252    pub fn with_ibct_ttl(mut self, ttl: Duration) -> Self {
253        self.ibct_ttl = ttl;
254        self
255    }
256
257    /// Issues and base64-encodes an `X-Zeph-IBCT` header value scoped to `endpoint` +
258    /// `task_id`. Returns `None` when no key is configured, or when issuance/encoding
259    /// fails (logged as a warning).
260    ///
261    /// `endpoint` is normalized to its origin (`scheme://host[:port]`, no path) before
262    /// being embedded in the token — see [`ibct_scope_origin`] for why: the server verifies
263    /// against its own advertised `AgentCard::url`, which is documented as the base URL with
264    /// no path suffix, and is shared by both the `/a2a` and `/a2a/stream` routes. Scoping the
265    /// token to the full request URL (including the route-specific path) would make it valid
266    /// for at most one of the two routes even against the correct server (#6260 review S1).
267    fn ibct_header_value(&self, endpoint: &str, task_id: &str) -> Option<String> {
268        let key = self.ibct_key.as_ref()?;
269        let scope = ibct_scope_origin(endpoint);
270        let token = match Ibct::issue(task_id, &scope, self.ibct_ttl, key) {
271            Ok(t) => t,
272            Err(e) => {
273                tracing::warn!("failed to issue IBCT token: {e}");
274                return None;
275            }
276        };
277        match token.encode() {
278            Ok(encoded) => Some(encoded),
279            Err(e) => {
280                tracing::warn!("failed to encode IBCT token: {e}");
281                None
282            }
283        }
284    }
285
286    /// # Errors
287    /// Returns `A2aError` on network, JSON, or JSON-RPC errors, or `A2aError::Timeout`
288    /// if the request exceeds the configured `request_timeout`.
289    #[tracing::instrument(name = "a2a.client.send_message", skip_all, err)]
290    pub async fn send_message(
291        &self,
292        endpoint: &str,
293        params: SendMessageParams,
294        token: Option<&str>,
295    ) -> Result<Task, A2aError> {
296        let task_id = params.message.task_id.clone().unwrap_or_default();
297        self.rpc_call(endpoint, METHOD_SEND_MESSAGE, params, token, &task_id)
298            .await
299    }
300
301    /// # Errors
302    /// Returns `A2aError` on network failure or if the SSE connection cannot be established.
303    #[tracing::instrument(name = "a2a.client.stream_message", skip_all, err)]
304    pub async fn stream_message(
305        &self,
306        endpoint: &str,
307        params: SendMessageParams,
308        token: Option<&str>,
309    ) -> Result<TaskEventStream, A2aError> {
310        let task_id = params.message.task_id.clone().unwrap_or_default();
311        let pinned = self.validate_endpoint(endpoint).await?;
312        let request_client = self.request_client(pinned.as_ref())?;
313        let request = JsonRpcRequest::new(METHOD_SEND_STREAMING_MESSAGE, params);
314        let mut req = request_client.post(endpoint).json(&request);
315        if let Some(t) = token {
316            req = req.bearer_auth(t);
317        }
318        if let Some(ibct_header) = self.ibct_header_value(endpoint, &task_id) {
319            req = req.header("X-Zeph-IBCT", ibct_header);
320        }
321        let resp = tokio::time::timeout(self.request_timeout, req.send())
322            .await
323            .map_err(|_| A2aError::Timeout(self.request_timeout))?
324            .map_err(A2aError::Http)?;
325
326        if !resp.status().is_success() {
327            let status = resp.status();
328            let body = tokio::time::timeout(Duration::from_secs(5), resp.text())
329                .await
330                .unwrap_or(Ok(String::new()))
331                .unwrap_or_default();
332            // Truncate body to avoid leaking large upstream error responses.
333            let truncated = if body.len() > 256 {
334                format!("{}…", &body[..256])
335            } else {
336                body
337            };
338            return Err(A2aError::Stream(format!("HTTP {status}: {truncated}")));
339        }
340
341        let event_stream = resp.bytes_stream().eventsource();
342        let mapped = event_stream.filter_map(|event| match event {
343            Ok(event) => {
344                if event.data.is_empty() || event.data == "[DONE]" {
345                    return None;
346                }
347                match serde_json::from_str::<JsonRpcResponse<TaskEvent>>(&event.data) {
348                    Ok(rpc_resp) => match rpc_resp.into_result() {
349                        Ok(task_event) => Some(Ok(task_event)),
350                        Err(rpc_err) => Some(Err(A2aError::from(rpc_err))),
351                    },
352                    Err(e) => Some(Err(A2aError::Stream(format!(
353                        "failed to parse SSE event: {e}"
354                    )))),
355                }
356            }
357            Err(e) => Some(Err(A2aError::Stream(format!("SSE stream error: {e}")))),
358        });
359
360        Ok(Box::pin(mapped))
361    }
362
363    /// # Errors
364    /// Returns `A2aError` on network, JSON, or JSON-RPC errors, or `A2aError::Timeout`
365    /// if the request exceeds the configured `request_timeout`.
366    #[tracing::instrument(name = "a2a.client.get_task", skip_all, err)]
367    pub async fn get_task(
368        &self,
369        endpoint: &str,
370        params: TaskIdParams,
371        token: Option<&str>,
372    ) -> Result<Task, A2aError> {
373        let task_id = params.id.clone();
374        self.rpc_call(endpoint, METHOD_GET_TASK, params, token, &task_id)
375            .await
376    }
377
378    /// # Errors
379    /// Returns `A2aError` on network, JSON, or JSON-RPC errors, or `A2aError::Timeout`
380    /// if the request exceeds the configured `request_timeout`.
381    #[tracing::instrument(name = "a2a.client.cancel_task", skip_all, err)]
382    pub async fn cancel_task(
383        &self,
384        endpoint: &str,
385        params: TaskIdParams,
386        token: Option<&str>,
387    ) -> Result<Task, A2aError> {
388        let task_id = params.id.clone();
389        self.rpc_call(endpoint, METHOD_CANCEL_TASK, params, token, &task_id)
390            .await
391    }
392
393    /// Validates `endpoint` against the configured [`SecurityPolicy`] and, when
394    /// `ssrf_protection` is enabled, resolves its hostname once and returns the
395    /// validated addresses to be pinned for the actual connection.
396    ///
397    /// Returning `Ok(None)` means either security is off for that check, or the
398    /// endpoint has no host (validation is skipped, matching prior behavior).
399    #[tracing::instrument(name = "a2a.client.validate_endpoint", skip_all, err)]
400    async fn validate_endpoint(&self, endpoint: &str) -> Result<Option<PinnedTarget>, A2aError> {
401        if self.security.require_tls && !endpoint.starts_with("https://") {
402            return Err(A2aError::Security(format!(
403                "TLS required but endpoint uses HTTP: {endpoint}"
404            )));
405        }
406
407        if !self.security.ssrf_protection {
408            return Ok(None);
409        }
410
411        let url: url::Url = endpoint
412            .parse()
413            .map_err(|e| A2aError::Security(format!("invalid URL: {e}")))?;
414
415        let Some(host) = url.host_str() else {
416            return Ok(None);
417        };
418        let port = url.port_or_known_default().unwrap_or(443);
419        let addrs = resolve_and_validate(host, port)
420            .await
421            .map_err(|e| A2aError::Security(e.to_string()))?;
422
423        Ok(Some(PinnedTarget {
424            host: host.to_owned(),
425            addrs,
426        }))
427    }
428
429    /// Returns `true` when either half of the [`SecurityPolicy`] requires requests
430    /// to be sent through a dedicated per-request client instead of `self.client`.
431    fn needs_hardened_client(&self) -> bool {
432        self.security.require_tls || self.security.ssrf_protection
433    }
434
435    /// Selects the `reqwest::Client` to use for a single request: the shared
436    /// injected client when no security is configured, or a fresh hardened client
437    /// (redirects disabled, optionally TLS-enforced and address-pinned) otherwise.
438    fn request_client(&self, pinned: Option<&PinnedTarget>) -> Result<reqwest::Client, A2aError> {
439        if self.needs_hardened_client() {
440            self.build_hardened_client(pinned)
441        } else {
442            Ok(self.client.clone())
443        }
444    }
445
446    /// Builds a per-request client hardened per the configured [`SecurityPolicy`].
447    ///
448    /// Redirects are always disabled (`Policy::none()`) so a malicious `3xx` response
449    /// cannot silently redirect the connection to a private address or downgrade to
450    /// `http://` — the caller (`rpc_call`/`stream_message`) treats any non-2xx or
451    /// unparseable response as an error instead of following it. When `pinned` is
452    /// `Some`, the client is additionally locked to the exact addresses that passed
453    /// SSRF validation via `resolve_to_addrs`, so reqwest cannot re-resolve the
454    /// hostname to a different address at connect time (closing the DNS-rebinding
455    /// TOCTOU window).
456    fn build_hardened_client(
457        &self,
458        pinned: Option<&PinnedTarget>,
459    ) -> Result<reqwest::Client, A2aError> {
460        let mut builder = reqwest::Client::builder()
461            .user_agent(concat!("zeph-a2a/", env!("CARGO_PKG_VERSION")))
462            .redirect(reqwest::redirect::Policy::none());
463
464        if self.security.require_tls {
465            builder = builder.https_only(true);
466        }
467        if let Some(target) = pinned {
468            builder = builder.resolve_to_addrs(&target.host, &target.addrs);
469        }
470
471        builder
472            .build()
473            .map_err(|e| A2aError::Security(format!("failed to build hardened client: {e}")))
474    }
475
476    #[tracing::instrument(name = "a2a.client.rpc_call", skip_all, err)]
477    async fn rpc_call<P: Serialize, R: DeserializeOwned>(
478        &self,
479        endpoint: &str,
480        method: &str,
481        params: P,
482        token: Option<&str>,
483        task_id: &str,
484    ) -> Result<R, A2aError> {
485        let pinned = self.validate_endpoint(endpoint).await?;
486        let request_client = self.request_client(pinned.as_ref())?;
487        let request = JsonRpcRequest::new(method, params);
488        let mut req = request_client.post(endpoint).json(&request);
489        if let Some(t) = token {
490            req = req.bearer_auth(t);
491        }
492        if let Some(ibct_header) = self.ibct_header_value(endpoint, task_id) {
493            req = req.header("X-Zeph-IBCT", ibct_header);
494        }
495        let rpc_response: JsonRpcResponse<R> = tokio::time::timeout(self.request_timeout, async {
496            let resp = req.send().await?;
497            resp.json().await
498        })
499        .await
500        .map_err(|_| A2aError::Timeout(self.request_timeout))?
501        .map_err(A2aError::Http)?;
502        rpc_response.into_result().map_err(A2aError::from)
503    }
504}
505
506#[cfg(test)]
507mod tests {
508    use std::assert_matches;
509    use std::net::IpAddr;
510
511    use super::*;
512    use zeph_common::net::is_private_ip;
513
514    use crate::jsonrpc::{JsonRpcError, JsonRpcResponse};
515    use crate::types::{
516        Artifact, Message, Part, Task, TaskArtifactUpdateEvent, TaskState, TaskStatus,
517        TaskStatusUpdateEvent,
518    };
519
520    #[test]
521    fn task_event_deserialize_status_update() {
522        let event = TaskStatusUpdateEvent {
523            kind: "status-update".into(),
524            task_id: "t-1".into(),
525            context_id: None,
526            status: TaskStatus {
527                state: TaskState::Working,
528                timestamp: "ts".into(),
529                message: Some(Message::user_text("thinking...")),
530            },
531            is_final: false,
532        };
533        let json = serde_json::to_string(&event).unwrap();
534        let parsed: TaskEvent = serde_json::from_str(&json).unwrap();
535        assert_matches!(parsed, TaskEvent::StatusUpdate(_));
536    }
537
538    #[test]
539    fn task_event_deserialize_artifact_update() {
540        let event = TaskArtifactUpdateEvent {
541            kind: "artifact-update".into(),
542            task_id: "t-1".into(),
543            context_id: None,
544            artifact: Artifact {
545                artifact_id: "a-1".into(),
546                name: None,
547                parts: vec![Part::text("result")],
548                metadata: None,
549            },
550            is_final: true,
551        };
552        let json = serde_json::to_string(&event).unwrap();
553        let parsed: TaskEvent = serde_json::from_str(&json).unwrap();
554        assert_matches!(parsed, TaskEvent::ArtifactUpdate(_));
555    }
556
557    #[test]
558    fn rpc_response_with_task_result() {
559        let task = Task {
560            id: "t-1".into(),
561            context_id: None,
562            status: TaskStatus {
563                state: TaskState::Completed,
564                timestamp: "ts".into(),
565                message: None,
566            },
567            artifacts: vec![],
568            history: vec![],
569            metadata: None,
570        };
571        let resp = JsonRpcResponse {
572            jsonrpc: "2.0".into(),
573            id: serde_json::Value::String("req-1".into()),
574            result: Some(task),
575            error: None,
576        };
577        let json = serde_json::to_string(&resp).unwrap();
578        let back: JsonRpcResponse<Task> = serde_json::from_str(&json).unwrap();
579        let task = back.into_result().unwrap();
580        assert_eq!(task.id, "t-1");
581        assert_eq!(task.status.state, TaskState::Completed);
582    }
583
584    #[test]
585    fn rpc_response_with_error() {
586        let resp: JsonRpcResponse<Task> = JsonRpcResponse {
587            jsonrpc: "2.0".into(),
588            id: serde_json::Value::String("req-1".into()),
589            result: None,
590            error: Some(JsonRpcError {
591                code: -32001,
592                message: "task not found".into(),
593                data: None,
594            }),
595        };
596        let json = serde_json::to_string(&resp).unwrap();
597        let back: JsonRpcResponse<Task> = serde_json::from_str(&json).unwrap();
598        let err = back.into_result().unwrap_err();
599        assert_eq!(err.code, -32001);
600    }
601
602    #[test]
603    fn a2a_client_construction() {
604        let client = A2aClient::new(reqwest::Client::new());
605        drop(client);
606    }
607
608    #[test]
609    fn is_private_ip_loopback() {
610        assert!(is_private_ip(IpAddr::V4(std::net::Ipv4Addr::LOCALHOST)));
611        assert!(is_private_ip(IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)));
612    }
613
614    #[test]
615    fn is_private_ip_private_ranges() {
616        assert!(is_private_ip("10.0.0.1".parse().unwrap()));
617        assert!(is_private_ip("172.16.0.1".parse().unwrap()));
618        assert!(is_private_ip("192.168.1.1".parse().unwrap()));
619    }
620
621    #[test]
622    fn is_private_ip_link_local() {
623        assert!(is_private_ip("169.254.0.1".parse().unwrap()));
624    }
625
626    #[test]
627    fn is_private_ip_unspecified() {
628        assert!(is_private_ip("0.0.0.0".parse().unwrap()));
629        assert!(is_private_ip("::".parse().unwrap()));
630    }
631
632    #[test]
633    fn is_private_ip_public() {
634        assert!(!is_private_ip("8.8.8.8".parse().unwrap()));
635        assert!(!is_private_ip("1.1.1.1".parse().unwrap()));
636    }
637
638    #[tokio::test]
639    async fn tls_enforcement_rejects_http() {
640        let client = A2aClient::new(reqwest::Client::new()).with_security(SecurityPolicy {
641            require_tls: true,
642            ssrf_protection: false,
643        });
644        let result = client.validate_endpoint("http://example.com/rpc").await;
645        assert!(result.is_err());
646        let err = result.unwrap_err();
647        assert_matches!(err, A2aError::Security(_));
648        assert!(err.to_string().contains("TLS required"));
649    }
650
651    #[tokio::test]
652    async fn tls_enforcement_allows_https() {
653        let client = A2aClient::new(reqwest::Client::new()).with_security(SecurityPolicy {
654            require_tls: true,
655            ssrf_protection: false,
656        });
657        let result = client.validate_endpoint("https://example.com/rpc").await;
658        assert!(result.is_ok());
659    }
660
661    #[tokio::test]
662    async fn ssrf_protection_rejects_localhost() {
663        let client = A2aClient::new(reqwest::Client::new()).with_security(SecurityPolicy {
664            require_tls: false,
665            ssrf_protection: true,
666        });
667        let result = client.validate_endpoint("http://127.0.0.1:8080/rpc").await;
668        assert!(result.is_err());
669        assert!(result.unwrap_err().to_string().contains("SSRF"));
670    }
671
672    #[tokio::test]
673    async fn no_security_allows_http_localhost() {
674        let client = A2aClient::new(reqwest::Client::new());
675        let result = client.validate_endpoint("http://127.0.0.1:8080/rpc").await;
676        assert!(result.is_ok());
677    }
678
679    #[test]
680    fn jsonrpc_request_serialization_for_send_message() {
681        let params = SendMessageParams {
682            message: Message::user_text("hello"),
683            configuration: None,
684        };
685        let req = JsonRpcRequest::new(METHOD_SEND_MESSAGE, params);
686        let json = serde_json::to_string(&req).unwrap();
687        assert!(json.contains("\"method\":\"message/send\""));
688        assert!(json.contains("\"jsonrpc\":\"2.0\""));
689        assert!(json.contains("\"hello\""));
690    }
691
692    #[test]
693    fn jsonrpc_request_serialization_for_get_task() {
694        let params = TaskIdParams {
695            id: "task-123".into(),
696            history_length: Some(5),
697        };
698        let req = JsonRpcRequest::new(METHOD_GET_TASK, params);
699        let json = serde_json::to_string(&req).unwrap();
700        assert!(json.contains("\"method\":\"tasks/get\""));
701        assert!(json.contains("\"task-123\""));
702        assert!(json.contains("\"historyLength\":5"));
703    }
704
705    #[test]
706    fn jsonrpc_request_serialization_for_cancel_task() {
707        let params = TaskIdParams {
708            id: "task-456".into(),
709            history_length: None,
710        };
711        let req = JsonRpcRequest::new(METHOD_CANCEL_TASK, params);
712        let json = serde_json::to_string(&req).unwrap();
713        assert!(json.contains("\"method\":\"tasks/cancel\""));
714        assert!(!json.contains("historyLength"));
715    }
716
717    #[test]
718    fn jsonrpc_request_serialization_for_stream() {
719        let params = SendMessageParams {
720            message: Message::user_text("stream me"),
721            configuration: None,
722        };
723        let req = JsonRpcRequest::new(METHOD_SEND_STREAMING_MESSAGE, params);
724        let json = serde_json::to_string(&req).unwrap();
725        assert!(json.contains("\"method\":\"message/stream\""));
726    }
727
728    #[tokio::test]
729    async fn send_message_connection_error() {
730        let client = A2aClient::new(reqwest::Client::new());
731        let params = SendMessageParams {
732            message: Message::user_text("hello"),
733            configuration: None,
734        };
735        let result = client
736            .send_message("http://127.0.0.1:1/rpc", params, None)
737            .await;
738        assert!(result.is_err());
739        assert_matches!(result.unwrap_err(), A2aError::Http(_));
740    }
741
742    #[tokio::test]
743    async fn get_task_connection_error() {
744        let client = A2aClient::new(reqwest::Client::new());
745        let params = TaskIdParams {
746            id: "t-1".into(),
747            history_length: None,
748        };
749        let result = client
750            .get_task("http://127.0.0.1:1/rpc", params, None)
751            .await;
752        assert!(result.is_err());
753        assert_matches!(result.unwrap_err(), A2aError::Http(_));
754    }
755
756    #[tokio::test]
757    async fn cancel_task_connection_error() {
758        let client = A2aClient::new(reqwest::Client::new());
759        let params = TaskIdParams {
760            id: "t-1".into(),
761            history_length: None,
762        };
763        let result = client
764            .cancel_task("http://127.0.0.1:1/rpc", params, None)
765            .await;
766        assert!(result.is_err());
767        assert_matches!(result.unwrap_err(), A2aError::Http(_));
768    }
769
770    #[tokio::test]
771    async fn stream_message_connection_error() {
772        let client = A2aClient::new(reqwest::Client::new());
773        let params = SendMessageParams {
774            message: Message::user_text("stream me"),
775            configuration: None,
776        };
777        let result = client
778            .stream_message("http://127.0.0.1:1/rpc", params, None)
779            .await;
780        assert!(result.is_err());
781    }
782
783    #[tokio::test]
784    async fn stream_message_tls_required_rejects_http() {
785        let client = A2aClient::new(reqwest::Client::new()).with_security(SecurityPolicy {
786            require_tls: true,
787            ssrf_protection: false,
788        });
789        let params = SendMessageParams {
790            message: Message::user_text("hello"),
791            configuration: None,
792        };
793        let result = client
794            .stream_message("http://example.com/rpc", params, None)
795            .await;
796        match result {
797            Err(A2aError::Security(msg)) => assert!(msg.contains("TLS required")),
798            _ => panic!("expected Security error"),
799        }
800    }
801
802    #[tokio::test]
803    async fn send_message_tls_required_rejects_http() {
804        let client = A2aClient::new(reqwest::Client::new()).with_security(SecurityPolicy {
805            require_tls: true,
806            ssrf_protection: false,
807        });
808        let params = SendMessageParams {
809            message: Message::user_text("hello"),
810            configuration: None,
811        };
812        let result = client
813            .send_message("http://example.com/rpc", params, None)
814            .await;
815        assert!(result.is_err());
816        assert_matches!(result.unwrap_err(), A2aError::Security(_));
817    }
818
819    #[tokio::test]
820    async fn get_task_tls_required_rejects_http() {
821        let client = A2aClient::new(reqwest::Client::new()).with_security(SecurityPolicy {
822            require_tls: true,
823            ssrf_protection: false,
824        });
825        let params = TaskIdParams {
826            id: "t-1".into(),
827            history_length: None,
828        };
829        let result = client
830            .get_task("http://example.com/rpc", params, None)
831            .await;
832        assert!(result.is_err());
833        assert_matches!(result.unwrap_err(), A2aError::Security(_));
834    }
835
836    #[tokio::test]
837    async fn cancel_task_tls_required_rejects_http() {
838        let client = A2aClient::new(reqwest::Client::new()).with_security(SecurityPolicy {
839            require_tls: true,
840            ssrf_protection: false,
841        });
842        let params = TaskIdParams {
843            id: "t-1".into(),
844            history_length: None,
845        };
846        let result = client
847            .cancel_task("http://example.com/rpc", params, None)
848            .await;
849        assert!(result.is_err());
850        assert_matches!(result.unwrap_err(), A2aError::Security(_));
851    }
852
853    #[tokio::test]
854    async fn validate_endpoint_invalid_url_with_ssrf() {
855        let client = A2aClient::new(reqwest::Client::new()).with_security(SecurityPolicy {
856            require_tls: false,
857            ssrf_protection: true,
858        });
859        let result = client.validate_endpoint("not-a-url").await;
860        assert!(result.is_err());
861        assert_matches!(result.unwrap_err(), A2aError::Security(_));
862    }
863
864    #[test]
865    fn with_security_returns_configured_client() {
866        let client =
867            A2aClient::new(reqwest::Client::new()).with_security(SecurityPolicy::hardened());
868        assert!(client.security.require_tls);
869        assert!(client.security.ssrf_protection);
870    }
871
872    #[test]
873    fn default_client_no_security() {
874        let client = A2aClient::new(reqwest::Client::new());
875        assert!(!client.security.require_tls);
876        assert!(!client.security.ssrf_protection);
877    }
878
879    #[test]
880    fn needs_hardened_client_reflects_policy() {
881        assert!(!A2aClient::new(reqwest::Client::new()).needs_hardened_client());
882        assert!(
883            A2aClient::new(reqwest::Client::new())
884                .with_security(SecurityPolicy {
885                    require_tls: true,
886                    ssrf_protection: false,
887                })
888                .needs_hardened_client()
889        );
890        assert!(
891            A2aClient::new(reqwest::Client::new())
892                .with_security(SecurityPolicy {
893                    require_tls: false,
894                    ssrf_protection: true,
895                })
896                .needs_hardened_client()
897        );
898        assert!(
899            A2aClient::new(reqwest::Client::new())
900                .with_security(SecurityPolicy::hardened())
901                .needs_hardened_client()
902        );
903    }
904
905    #[test]
906    fn task_event_clone() {
907        let event = TaskEvent::StatusUpdate(TaskStatusUpdateEvent {
908            kind: "status-update".into(),
909            task_id: "t-1".into(),
910            context_id: None,
911            status: TaskStatus {
912                state: TaskState::Working,
913                timestamp: "ts".into(),
914                message: None,
915            },
916            is_final: false,
917        });
918        let cloned = event.clone();
919        let json1 = serde_json::to_string(&event).unwrap();
920        let json2 = serde_json::to_string(&cloned).unwrap();
921        assert_eq!(json1, json2);
922    }
923
924    #[test]
925    fn task_event_debug() {
926        let event = TaskEvent::ArtifactUpdate(TaskArtifactUpdateEvent {
927            kind: "artifact-update".into(),
928            task_id: "t-1".into(),
929            context_id: None,
930            artifact: Artifact {
931                artifact_id: "a-1".into(),
932                name: None,
933                parts: vec![Part::text("data")],
934                metadata: None,
935            },
936            is_final: true,
937        });
938        let dbg = format!("{event:?}");
939        assert!(dbg.contains("ArtifactUpdate"));
940    }
941
942    #[test]
943    fn is_private_ip_ipv4_non_private() {
944        assert!(!is_private_ip("93.184.216.34".parse().unwrap()));
945    }
946
947    #[test]
948    fn is_private_ip_ipv6_non_private() {
949        assert!(!is_private_ip("2001:db8::1".parse().unwrap()));
950    }
951
952    #[test]
953    fn rpc_response_error_takes_priority_over_result() {
954        let resp = JsonRpcResponse {
955            jsonrpc: "2.0".into(),
956            id: serde_json::Value::String("1".into()),
957            result: Some(Task {
958                id: "t-1".into(),
959                context_id: None,
960                status: TaskStatus {
961                    state: TaskState::Completed,
962                    timestamp: "ts".into(),
963                    message: None,
964                },
965                artifacts: vec![],
966                history: vec![],
967                metadata: None,
968            }),
969            error: Some(JsonRpcError {
970                code: -32001,
971                message: "error".into(),
972                data: None,
973            }),
974        };
975        let err = resp.into_result().unwrap_err();
976        assert_eq!(err.code, -32001);
977    }
978
979    #[test]
980    fn rpc_response_neither_result_nor_error() {
981        let resp: JsonRpcResponse<Task> = JsonRpcResponse {
982            jsonrpc: "2.0".into(),
983            id: serde_json::Value::String("1".into()),
984            result: None,
985            error: None,
986        };
987        let err = resp.into_result().unwrap_err();
988        assert_eq!(err.code, -32603);
989    }
990
991    #[test]
992    fn with_ibct_key_sets_key_and_default_ttl() {
993        let key = IbctKey {
994            key_id: "k1".into(),
995            key_bytes: b"secret".to_vec(),
996        };
997        let client = A2aClient::new(reqwest::Client::new()).with_ibct_key(key);
998        assert!(client.ibct_key.is_some());
999        assert_eq!(client.ibct_ttl, Duration::from_mins(5));
1000    }
1001
1002    #[test]
1003    fn with_ibct_ttl_overrides_default() {
1004        let key = IbctKey {
1005            key_id: "k1".into(),
1006            key_bytes: b"secret".to_vec(),
1007        };
1008        let client = A2aClient::new(reqwest::Client::new())
1009            .with_ibct_key(key)
1010            .with_ibct_ttl(Duration::from_mins(1));
1011        assert_eq!(client.ibct_ttl, Duration::from_mins(1));
1012    }
1013
1014    #[test]
1015    fn no_ibct_key_configured_yields_no_header() {
1016        let client = A2aClient::new(reqwest::Client::new());
1017        assert!(
1018            client
1019                .ibct_header_value("https://agent.example.com", "task-1")
1020                .is_none()
1021        );
1022    }
1023
1024    #[cfg(feature = "ibct")]
1025    #[test]
1026    fn ibct_key_configured_yields_encoded_header() {
1027        let key = IbctKey {
1028            key_id: "k1".into(),
1029            key_bytes: b"secret".to_vec(),
1030        };
1031        let client = A2aClient::new(reqwest::Client::new()).with_ibct_key(key);
1032        let header = client
1033            .ibct_header_value("https://agent.example.com", "task-1")
1034            .expect("header should be issued when ibct feature is enabled");
1035        let decoded = crate::ibct::Ibct::decode(&header).unwrap();
1036        assert_eq!(decoded.task_id, "task-1");
1037        assert_eq!(decoded.endpoint, "https://agent.example.com");
1038    }
1039
1040    // Regression tests for #6260 review S1: the server verifies every IBCT against its own
1041    // pathless `AgentCard::url`, shared by both the `/a2a` and `/a2a/stream` routes. A token
1042    // scoped to the full per-route POST URL (including the route-specific path) could match
1043    // at most one of the two routes even against the correct server. `ibct_scope_origin` (and
1044    // therefore `ibct_header_value`) must always strip the path before issuing.
1045
1046    #[test]
1047    fn ibct_scope_origin_strips_path_and_query() {
1048        assert_eq!(
1049            ibct_scope_origin("https://agent.example.com/a2a"),
1050            "https://agent.example.com"
1051        );
1052        assert_eq!(
1053            ibct_scope_origin("https://agent.example.com/a2a/stream"),
1054            "https://agent.example.com"
1055        );
1056        assert_eq!(
1057            ibct_scope_origin("http://127.0.0.1:8080/a2a/stream?x=1"),
1058            "http://127.0.0.1:8080"
1059        );
1060    }
1061
1062    #[test]
1063    fn ibct_scope_origin_agrees_across_both_a2a_routes() {
1064        let a2a = ibct_scope_origin("http://127.0.0.1:8080/a2a");
1065        let stream = ibct_scope_origin("http://127.0.0.1:8080/a2a/stream");
1066        assert_eq!(
1067            a2a, stream,
1068            "both routes on the same agent must scope to the same IBCT endpoint"
1069        );
1070    }
1071
1072    #[test]
1073    fn ibct_scope_origin_falls_back_to_input_on_unparseable_url() {
1074        assert_eq!(ibct_scope_origin("not-a-url"), "not-a-url");
1075    }
1076
1077    #[cfg(feature = "ibct")]
1078    #[test]
1079    fn ibct_header_value_scopes_token_to_origin_not_full_path() {
1080        let key = IbctKey {
1081            key_id: "k1".into(),
1082            key_bytes: b"secret".to_vec(),
1083        };
1084        let client = A2aClient::new(reqwest::Client::new()).with_ibct_key(key);
1085
1086        let header_a2a = client
1087            .ibct_header_value("http://127.0.0.1:8080/a2a", "task-1")
1088            .unwrap();
1089        let header_stream = client
1090            .ibct_header_value("http://127.0.0.1:8080/a2a/stream", "task-1")
1091            .unwrap();
1092
1093        let decoded_a2a = crate::ibct::Ibct::decode(&header_a2a).unwrap();
1094        let decoded_stream = crate::ibct::Ibct::decode(&header_stream).unwrap();
1095        assert_eq!(decoded_a2a.endpoint, "http://127.0.0.1:8080");
1096        assert_eq!(
1097            decoded_a2a.endpoint, decoded_stream.endpoint,
1098            "a token issued for /a2a and one issued for /a2a/stream on the same agent must \
1099             carry the same endpoint scope, since the server verifies both against one \
1100             pathless card.url"
1101        );
1102    }
1103
1104    #[test]
1105    fn task_event_serialize_round_trip() {
1106        let event = TaskEvent::StatusUpdate(TaskStatusUpdateEvent {
1107            kind: "status-update".into(),
1108            task_id: "t-1".into(),
1109            context_id: Some("ctx-1".into()),
1110            status: TaskStatus {
1111                state: TaskState::Completed,
1112                timestamp: "2025-01-01T00:00:00Z".into(),
1113                message: Some(Message::user_text("done")),
1114            },
1115            is_final: true,
1116        });
1117        let json = serde_json::to_string(&event).unwrap();
1118        let back: TaskEvent = serde_json::from_str(&json).unwrap();
1119        assert_matches!(back, TaskEvent::StatusUpdate(_));
1120    }
1121}
1122
1123#[cfg(test)]
1124mod wiremock_tests {
1125    use std::assert_matches;
1126    use tokio_stream::StreamExt;
1127    use wiremock::matchers::{header, method, path};
1128    use wiremock::{Mock, MockServer, ResponseTemplate};
1129
1130    use crate::client::{A2aClient, PinnedTarget, SecurityPolicy};
1131    use crate::jsonrpc::{SendMessageParams, TaskIdParams};
1132    use crate::testing::*;
1133    use crate::types::Message;
1134
1135    #[tokio::test]
1136    async fn send_message_success() {
1137        let server = MockServer::start().await;
1138        Mock::given(method("POST"))
1139            .and(path("/rpc"))
1140            .respond_with(task_rpc_response("task-1", "submitted"))
1141            .mount(&server)
1142            .await;
1143
1144        let client = A2aClient::new(reqwest::Client::new());
1145        let params = SendMessageParams {
1146            message: Message::user_text("hello"),
1147            configuration: None,
1148        };
1149        let task = client
1150            .send_message(&format!("{}/rpc", server.uri()), params, None)
1151            .await
1152            .unwrap();
1153        assert_eq!(task.id, "task-1");
1154    }
1155
1156    #[tokio::test]
1157    async fn send_message_rpc_error() {
1158        let server = MockServer::start().await;
1159        Mock::given(method("POST"))
1160            .and(path("/rpc"))
1161            .respond_with(task_rpc_error_response(-32001, "task not found"))
1162            .mount(&server)
1163            .await;
1164
1165        let client = A2aClient::new(reqwest::Client::new());
1166        let params = SendMessageParams {
1167            message: Message::user_text("hi"),
1168            configuration: None,
1169        };
1170        let result = client
1171            .send_message(&format!("{}/rpc", server.uri()), params, None)
1172            .await;
1173        assert!(result.is_err());
1174        let err = result.unwrap_err();
1175        assert_matches!(err, crate::error::A2aError::JsonRpc { code: -32001, .. });
1176    }
1177
1178    #[tokio::test]
1179    async fn send_message_with_bearer_auth() {
1180        let server = MockServer::start().await;
1181        Mock::given(method("POST"))
1182            .and(path("/rpc"))
1183            .and(header("authorization", "Bearer secret-token"))
1184            .respond_with(task_rpc_response("task-auth", "submitted"))
1185            .mount(&server)
1186            .await;
1187
1188        let client = A2aClient::new(reqwest::Client::new());
1189        let params = SendMessageParams {
1190            message: Message::user_text("secure"),
1191            configuration: None,
1192        };
1193        let task = client
1194            .send_message(
1195                &format!("{}/rpc", server.uri()),
1196                params,
1197                Some("secret-token"),
1198            )
1199            .await
1200            .unwrap();
1201        assert_eq!(task.id, "task-auth");
1202    }
1203
1204    #[tokio::test]
1205    async fn get_task_success() {
1206        let server = MockServer::start().await;
1207        Mock::given(method("POST"))
1208            .and(path("/rpc"))
1209            .respond_with(task_rpc_response("task-get", "completed"))
1210            .mount(&server)
1211            .await;
1212
1213        let client = A2aClient::new(reqwest::Client::new());
1214        let params = TaskIdParams {
1215            id: "task-get".into(),
1216            history_length: None,
1217        };
1218        let task = client
1219            .get_task(&format!("{}/rpc", server.uri()), params, None)
1220            .await
1221            .unwrap();
1222        assert_eq!(task.id, "task-get");
1223    }
1224
1225    #[tokio::test]
1226    async fn cancel_task_success() {
1227        let server = MockServer::start().await;
1228        Mock::given(method("POST"))
1229            .and(path("/rpc"))
1230            .respond_with(task_rpc_response("task-cancel", "canceled"))
1231            .mount(&server)
1232            .await;
1233
1234        let client = A2aClient::new(reqwest::Client::new());
1235        let params = TaskIdParams {
1236            id: "task-cancel".into(),
1237            history_length: None,
1238        };
1239        let task = client
1240            .cancel_task(&format!("{}/rpc", server.uri()), params, None)
1241            .await
1242            .unwrap();
1243        assert_eq!(task.id, "task-cancel");
1244    }
1245
1246    #[tokio::test]
1247    async fn stream_message_success() {
1248        let server = MockServer::start().await;
1249        Mock::given(method("POST"))
1250            .and(path("/rpc"))
1251            .respond_with(sse_task_events_response("task-stream", "result content"))
1252            .mount(&server)
1253            .await;
1254
1255        let client = A2aClient::new(reqwest::Client::new());
1256        let params = SendMessageParams {
1257            message: Message::user_text("stream"),
1258            configuration: None,
1259        };
1260        let stream = client
1261            .stream_message(&format!("{}/rpc", server.uri()), params, None)
1262            .await
1263            .unwrap();
1264        let events: Vec<_> = stream.collect().await;
1265        assert!(!events.is_empty());
1266    }
1267
1268    #[tokio::test]
1269    async fn stream_message_http_error() {
1270        let server = MockServer::start().await;
1271        Mock::given(method("POST"))
1272            .and(path("/rpc"))
1273            .respond_with(ResponseTemplate::new(500).set_body_string("Internal Server Error"))
1274            .mount(&server)
1275            .await;
1276
1277        let client = A2aClient::new(reqwest::Client::new());
1278        let params = SendMessageParams {
1279            message: Message::user_text("fail"),
1280            configuration: None,
1281        };
1282        let result = client
1283            .stream_message(&format!("{}/rpc", server.uri()), params, None)
1284            .await;
1285        let err = result.err().expect("expected error");
1286        assert_matches!(err, crate::error::A2aError::Stream(_));
1287    }
1288
1289    #[tokio::test]
1290    async fn rpc_call_times_out() {
1291        let server = MockServer::start().await;
1292        Mock::given(method("POST"))
1293            .and(path("/rpc"))
1294            .respond_with(
1295                ResponseTemplate::new(200)
1296                    .set_delay(std::time::Duration::from_secs(5))
1297                    .set_body_json(serde_json::json!({
1298                        "jsonrpc": "2.0",
1299                        "id": "req-1",
1300                        "result": {
1301                            "id": "t-1",
1302                            "status": {"state": "completed", "timestamp": "2026-01-01T00:00:00Z"}
1303                        }
1304                    })),
1305            )
1306            .mount(&server)
1307            .await;
1308
1309        let client = A2aClient::new(reqwest::Client::new())
1310            .with_request_timeout(std::time::Duration::from_millis(100));
1311        let params = SendMessageParams {
1312            message: Message::user_text("hello"),
1313            configuration: None,
1314        };
1315        let result = client
1316            .send_message(&format!("{}/rpc", server.uri()), params, None)
1317            .await;
1318        assert!(result.is_err());
1319        assert!(
1320            matches!(result.unwrap_err(), crate::error::A2aError::Timeout(_)),
1321            "expected Timeout error"
1322        );
1323    }
1324
1325    /// Proves the DNS-rebinding TOCTOU is closed: `resolve_to_addrs` pins the connection to
1326    /// the address validated by `resolve_and_validate`, so reqwest never re-resolves `fake_host`
1327    /// (a hostname reserved by RFC 2606 and guaranteed to never resolve via real DNS) at connect
1328    /// time. If the client re-resolved instead of using the pinned address, this request would
1329    /// fail with a DNS lookup error rather than reaching the mock server.
1330    #[tokio::test]
1331    async fn hardened_client_pins_connection_bypassing_dns() {
1332        let server = MockServer::start().await;
1333        Mock::given(method("GET"))
1334            .respond_with(ResponseTemplate::new(200).set_body_string("ok"))
1335            .mount(&server)
1336            .await;
1337
1338        let addr = *server.address();
1339        let fake_host = "zeph-a2a-pin-test.invalid";
1340        let client = A2aClient::new(reqwest::Client::new()).with_security(SecurityPolicy {
1341            require_tls: false,
1342            ssrf_protection: true,
1343        });
1344        let pinned = PinnedTarget {
1345            host: fake_host.to_owned(),
1346            addrs: vec![addr],
1347        };
1348        let hardened = client.build_hardened_client(Some(&pinned)).unwrap();
1349
1350        let resp = hardened
1351            .get(format!("http://{fake_host}:{}/", addr.port()))
1352            .send()
1353            .await
1354            .unwrap_or_else(|e| panic!("pinned request to unresolvable host failed: {e}"));
1355        assert_eq!(resp.status(), 200);
1356        assert_eq!(resp.text().await.unwrap(), "ok");
1357    }
1358
1359    /// Proves the redirect-based SSRF bypass is closed: the hardened client does not
1360    /// automatically follow a `3xx` response, even when `Location` points at a private
1361    /// address. `rpc_call`/`stream_message` treat the raw redirect response as a normal
1362    /// (non-2xx) response and surface an error instead of connecting to `Location`.
1363    #[tokio::test]
1364    async fn hardened_client_does_not_auto_follow_redirect_to_private_ip() {
1365        let server = MockServer::start().await;
1366        Mock::given(method("GET"))
1367            .respond_with(
1368                ResponseTemplate::new(302).insert_header("Location", "http://127.0.0.1:9/internal"),
1369            )
1370            .mount(&server)
1371            .await;
1372
1373        let addr = *server.address();
1374        let fake_host = "zeph-a2a-redirect-test.invalid";
1375        let client = A2aClient::new(reqwest::Client::new()).with_security(SecurityPolicy {
1376            require_tls: false,
1377            ssrf_protection: true,
1378        });
1379        let pinned = PinnedTarget {
1380            host: fake_host.to_owned(),
1381            addrs: vec![addr],
1382        };
1383        let hardened = client.build_hardened_client(Some(&pinned)).unwrap();
1384
1385        let resp = hardened
1386            .get(format!("http://{fake_host}:{}/start", addr.port()))
1387            .send()
1388            .await
1389            .unwrap();
1390
1391        assert_eq!(resp.status(), reqwest::StatusCode::FOUND);
1392        assert_eq!(
1393            resp.headers().get(reqwest::header::LOCATION).unwrap(),
1394            "http://127.0.0.1:9/internal"
1395        );
1396    }
1397
1398    /// Proves TLS enforcement holds even on the hardened per-request client: `https_only(true)`
1399    /// rejects a plaintext `http://` connection outright, closing the https-to-http downgrade
1400    /// gap that an unvalidated redirect could otherwise exploit.
1401    #[tokio::test]
1402    async fn hardened_client_with_require_tls_rejects_plaintext_connection() {
1403        let server = MockServer::start().await;
1404        Mock::given(method("GET"))
1405            .respond_with(ResponseTemplate::new(200))
1406            .mount(&server)
1407            .await;
1408
1409        let addr = *server.address();
1410        let fake_host = "zeph-a2a-tls-test.invalid";
1411        let client = A2aClient::new(reqwest::Client::new()).with_security(SecurityPolicy {
1412            require_tls: true,
1413            ssrf_protection: true,
1414        });
1415        let pinned = PinnedTarget {
1416            host: fake_host.to_owned(),
1417            addrs: vec![addr],
1418        };
1419        let hardened = client.build_hardened_client(Some(&pinned)).unwrap();
1420
1421        let result = hardened
1422            .get(format!("http://{fake_host}:{}/", addr.port()))
1423            .send()
1424            .await;
1425        assert!(
1426            result.is_err(),
1427            "https_only(true) must reject a plain http:// URL"
1428        );
1429    }
1430
1431    #[cfg(feature = "ibct")]
1432    #[tokio::test]
1433    async fn send_message_attaches_ibct_header_when_configured() {
1434        use wiremock::matchers::header_exists;
1435
1436        let server = MockServer::start().await;
1437        Mock::given(method("POST"))
1438            .and(path("/rpc"))
1439            .and(header_exists("x-zeph-ibct"))
1440            .respond_with(task_rpc_response("task-ibct", "submitted"))
1441            .mount(&server)
1442            .await;
1443
1444        let key = crate::ibct::IbctKey {
1445            key_id: "k1".into(),
1446            key_bytes: b"secret".to_vec(),
1447        };
1448        let client = A2aClient::new(reqwest::Client::new()).with_ibct_key(key);
1449        let params = SendMessageParams {
1450            message: Message::user_text("hello"),
1451            configuration: None,
1452        };
1453        let task = client
1454            .send_message(&format!("{}/rpc", server.uri()), params, None)
1455            .await
1456            .unwrap();
1457        assert_eq!(task.id, "task-ibct");
1458    }
1459}
zeph_a2a/client.rs

zeph_a2a/
client.rs
