ant_protocol/
chunk_protocol.rs

1//! Shared helper for the chunk protocol request/response pattern.
2//!
3//! Extracts the duplicated "subscribe → send → poll event loop" into a single
4//! generic function used by both `ant-client` and `ant-node` E2E helpers.
5
6use crate::chunk::{ChunkMessage, ChunkMessageBody, CHUNK_PROTOCOL_ID};
7use crate::logging::{debug, warn};
8use saorsa_core::identity::PeerId;
9use saorsa_core::{MultiAddr, P2PEvent, P2PNode};
10use std::time::Duration;
11use tokio::sync::broadcast::error::RecvError;
12use tokio::time::Instant;
13
14/// A decoded chunk-protocol response together with transport provenance.
15///
16/// `transport_source` is supplied by the authenticated transport receive path.
17/// It is diagnostic metadata, not an identity signal; `source_peer` remains
18/// the authenticated application-level peer identity used by the response
19/// filter.
20#[derive(Debug)]
21pub struct ChunkProtocolResponse<T, E> {
22    /// Result produced by the caller's response handler. Keeping the result
23    /// inside the envelope preserves provenance for structured remote errors.
24    pub result: Result<T, E>,
25    /// Authenticated peer that supplied the matching response.
26    pub source_peer: PeerId,
27    /// Transport address that delivered the response, when available.
28    pub transport_source: Option<MultiAddr>,
29}
30
31/// Send a chunk-protocol message to `target_peer` and await a matching response.
32///
33/// The event loop filters by topic (`CHUNK_PROTOCOL_ID`), source peer, decode
34/// errors (warn + skip), and `request_id` mismatch (skip).
35///
36/// * `response_handler` — inspects the decoded [`ChunkMessageBody`] and returns:
37///   - `Some(Ok(T))` to resolve successfully,
38///   - `Some(Err(E))` to resolve with an error,
39///   - `None` to keep waiting (wrong variant / not our response).
40/// * `send_error` — produces the caller's error type when `send_message` fails.
41/// * `timeout_error` — produces the caller's error type on deadline expiry.
42///
43/// # Errors
44///
45/// Returns `Err(E)` if sending fails (via `send_error`), the `response_handler`
46/// returns a protocol-level error, or the deadline expires (via `timeout_error`).
47#[allow(clippy::too_many_arguments)]
48pub async fn send_and_await_chunk_response<T, E>(
49    node: &P2PNode,
50    target_peer: &PeerId,
51    message_bytes: Vec<u8>,
52    request_id: u64,
53    timeout: Duration,
54    peer_addrs: &[MultiAddr],
55    response_handler: impl Fn(ChunkMessageBody) -> Option<Result<T, E>>,
56    send_error: impl FnOnce(String) -> E,
57    timeout_error: impl FnOnce() -> E,
58) -> Result<T, E> {
59    send_and_await_chunk_response_with_metadata(
60        node,
61        target_peer,
62        message_bytes,
63        request_id,
64        timeout,
65        peer_addrs,
66        response_handler,
67        send_error,
68        timeout_error,
69    )
70    .await
71    .and_then(|response| response.result)
72}
73
74/// Send a chunk-protocol message and return the decoded response plus the
75/// observed transport provenance.
76///
77/// This follows the same filtering and timeout behaviour as
78/// [`send_and_await_chunk_response`]. The additional metadata is captured from
79/// the already-received event and does not alter peer selection, dialing,
80/// retries, or response acceptance.
81///
82/// # Errors
83///
84/// Returns `Err(E)` if sending fails (via `send_error`) or the deadline
85/// expires (via `timeout_error`). A decoded protocol-level response error is
86/// retained in [`ChunkProtocolResponse::result`] so its source metadata is not
87/// lost.
88#[allow(clippy::too_many_arguments)]
89pub async fn send_and_await_chunk_response_with_metadata<T, E>(
90    node: &P2PNode,
91    target_peer: &PeerId,
92    message_bytes: Vec<u8>,
93    request_id: u64,
94    timeout: Duration,
95    peer_addrs: &[MultiAddr],
96    response_handler: impl Fn(ChunkMessageBody) -> Option<Result<T, E>>,
97    send_error: impl FnOnce(String) -> E,
98    timeout_error: impl FnOnce() -> E,
99) -> Result<ChunkProtocolResponse<T, E>, E> {
100    // Subscribe before sending so we don't miss the response
101    let mut events = node.subscribe_events();
102
103    node.send_message(target_peer, CHUNK_PROTOCOL_ID, message_bytes, peer_addrs)
104        .await
105        .map_err(|e| send_error(e.to_string()))?;
106
107    // `Instant::now() + timeout` can panic on extreme durations; fall back
108    // to the current instant (immediate timeout) if the addition overflows
109    // rather than taking down a crate that denies panics.
110    let deadline = Instant::now()
111        .checked_add(timeout)
112        .unwrap_or_else(Instant::now);
113
114    while Instant::now() < deadline {
115        let remaining = deadline.saturating_duration_since(Instant::now());
116        match tokio::time::timeout(remaining, events.recv()).await {
117            Ok(Ok(P2PEvent::Message {
118                topic,
119                source: Some(source),
120                transport_source,
121                data,
122                ..
123            })) if topic == CHUNK_PROTOCOL_ID && source == *target_peer => {
124                let response = match ChunkMessage::decode(&data) {
125                    Ok(r) => r,
126                    Err(e) => {
127                        warn!("Failed to decode chunk message, skipping: {e}");
128                        continue;
129                    }
130                };
131                if response.request_id != request_id {
132                    continue;
133                }
134                if let Some(result) = response_handler(response.body) {
135                    return Ok(ChunkProtocolResponse {
136                        result,
137                        source_peer: source,
138                        transport_source,
139                    });
140                }
141            }
142            Ok(Ok(_)) => {}
143            Ok(Err(RecvError::Lagged(skipped))) => {
144                debug!("Chunk protocol events lagged by {skipped} messages, continuing");
145            }
146            Ok(Err(RecvError::Closed)) | Err(_) => break,
147        }
148    }
149
150    Err(timeout_error())
151}
ant_protocol/chunk_protocol.rs

ant_protocol/
chunk_protocol.rs
