diff --git a/sdk/rust/src/lib.rs b/sdk/rust/src/lib.rs index 33a708b2..8e75f59e 100644 --- a/sdk/rust/src/lib.rs +++ b/sdk/rust/src/lib.rs @@ -55,7 +55,7 @@ pub use solana::{ build_exact_svm_transfer_transaction, derive_associated_token_address, get_recent_blockhash, ExactSvmTransfer, ExactSvmTransferOptions, }; -pub use websocket::{TinyPlaceWebSocket, WebSocketConnection, WsAuth}; +pub use websocket::{ReconnectPolicy, TinyPlaceWebSocket, WebSocketConnection, WsAuth}; pub use x402_standard::{ build_exact_svm_payment_payload, decode_payment_required, decode_settlement_response, encode_payment_signature, select_exact_svm_requirement, X402PaymentPayload, diff --git a/sdk/rust/src/websocket.rs b/sdk/rust/src/websocket.rs index 789cbb67..8cf2bec4 100644 --- a/sdk/rust/src/websocket.rs +++ b/sdk/rust/src/websocket.rs @@ -10,8 +10,16 @@ //! URL-borne credentials (directory-write query params, or the agent //! `Authorization` as an `authorization=` param) are still appended, so the same //! handle authenticates against backends that only read the query string. +//! +//! A dropped stream is re-established transparently by [`WebSocketConnection::recv`] +//! under a [`ReconnectPolicy`], whose defaults match the TypeScript SDK's +//! `reconnect` / `reconnectInterval` / `maxReconnectAttempts` options. Each +//! attempt re-runs the full dial, so the upgrade carries a **freshly signed** +//! credential — replaying the original request would present a stale +//! timestamp/nonce and be rejected with a 401. use std::sync::Arc; +use std::time::Duration; use futures_util::{SinkExt as _, StreamExt as _}; use tokio::net::TcpStream; @@ -40,14 +48,63 @@ pub enum WsAuth { Directory, } +/// How a stream that drops mid-flight is re-established. +/// +/// Defaults mirror the TypeScript SDK's `TinyPlaceWebSocketOptions`: reconnect +/// enabled, a 3s pause between attempts, and at most 10 attempts. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ReconnectPolicy { + /// Whether a dropped stream is re-dialled at all. + pub enabled: bool, + /// How long to wait before each attempt. + pub interval: Duration, + /// How many consecutive attempts to make before giving up. + /// + /// The counter resets as soon as a re-established stream carries *any* + /// frame — a keepalive ping counts — so this bounds a run of consecutive + /// failures rather than the lifetime total. A quiet stream still resets it; + /// only a socket that produces nothing at all before dropping again does + /// not. + pub max_attempts: u32, + /// How long a single re-dial may take before it counts as a failed attempt. + /// + /// `connect_async` has no timeout of its own, so a black-holed TCP connect + /// or a handshake the server never answers would otherwise hang `recv` + /// indefinitely and leave `max_attempts` bounding nothing. + pub connect_timeout: Duration, +} + +impl Default for ReconnectPolicy { + fn default() -> Self { + Self { + enabled: true, + interval: Duration::from_millis(3_000), + max_attempts: 10, + connect_timeout: Duration::from_secs(10), + } + } +} + +impl ReconnectPolicy { + /// A policy that never reconnects: `recv` reports the drop instead. + pub fn disabled() -> Self { + Self { + enabled: false, + ..Self::default() + } + } +} + /// An un-connected WebSocket handle. Cheap to build; opens the socket on /// [`connect`](TinyPlaceWebSocket::connect). +#[derive(Clone)] pub struct TinyPlaceWebSocket { origin: String, request_uri: String, signer: Option>, public_key: Option, auth: WsAuth, + reconnect: ReconnectPolicy, } impl TinyPlaceWebSocket { @@ -66,9 +123,31 @@ impl TinyPlaceWebSocket { signer, public_key, auth, + reconnect: ReconnectPolicy::default(), } } + /// Override how a dropped stream is re-established. + /// + /// ```no_run + /// # use tinyplace::{ReconnectPolicy, TinyPlaceClient}; + /// # async fn run(client: TinyPlaceClient) -> tinyplace::Result<()> { + /// // Give up immediately, surfacing the drop to the caller. + /// let mut conn = client + /// .inbox + /// .stream() + /// .with_reconnect(ReconnectPolicy::disabled()) + /// .connect() + /// .await?; + /// # let _ = conn.recv().await; + /// # Ok(()) + /// # } + /// ``` + pub fn with_reconnect(mut self, reconnect: ReconnectPolicy) -> Self { + self.reconnect = reconnect; + self + } + /// The fully-qualified WebSocket URL that would be opened, with auth applied. /// Exposed for testing and diagnostics. pub async fn signed_url(&self) -> Result { @@ -105,7 +184,27 @@ impl TinyPlaceWebSocket { } /// Open the WebSocket and return a connection to read/write messages. + /// + /// This is a single attempt: a refused or unauthenticated upgrade returns + /// the error rather than retrying, so a misconfigured client fails fast. + /// [`ReconnectPolicy`] governs what happens to an *established* stream that + /// later drops. pub async fn connect(&self) -> Result { + let stream = self.dial().await?; + Ok(WebSocketConnection { + inner: stream, + dialer: self.clone(), + attempts: 0, + last_dial_error: None, + }) + } + + /// Run one full dial: sign the URL and the upgrade headers, then handshake. + /// + /// Every reconnect goes through here rather than replaying a stored request: + /// [`sign_websocket_upgrade`] binds a timestamp and nonce, so a replayed + /// upgrade would be rejected as stale. + async fn dial(&self) -> Result>> { let url = self.signed_url().await?; let mut request = url .as_str() @@ -121,7 +220,7 @@ impl TinyPlaceWebSocket { let (stream, _response) = connect_async(request) .await .map_err(|error| Error::WebSocket(error.to_string()))?; - Ok(WebSocketConnection { inner: stream }) + Ok(stream) } } @@ -129,25 +228,112 @@ impl TinyPlaceWebSocket { /// push with [`send`](Self::send), and shut down with [`close`](Self::close). pub struct WebSocketConnection { inner: WebSocketStream>, + /// The handle this was opened from, kept so a drop can be re-dialled. + dialer: TinyPlaceWebSocket, + /// Consecutive reconnect attempts since the stream last carried a frame. + attempts: u32, + /// Why the most recent re-dial failed, reported in preference to the older + /// error that started the reconnect run. + last_dial_error: Option, } impl WebSocketConnection { - /// Await the next message, parsed as JSON. Returns `None` when the stream - /// closes. Ping/pong control frames are handled transparently and skipped. + /// Await the next message, parsed as JSON. Ping/pong control frames are + /// handled transparently and skipped. + /// + /// If the stream drops, it is re-dialled per the handle's + /// [`ReconnectPolicy`] and the wait resumes — transparently to the caller. + /// Returns `None` once the stream has closed and reconnecting is disabled or + /// exhausted. A transport error is likewise reported only after reconnecting + /// has been given up on. + /// + /// Note that a reconnect re-subscribes: anything the server published while + /// the socket was down is not replayed, and a stream that opens with a + /// snapshot frame will deliver a fresh one. pub async fn recv(&mut self) -> Option> { loop { match self.inner.next().await { Some(Ok(Message::Text(text))) => { + self.mark_healthy(); return Some(serde_json::from_str(&text).map_err(Error::from)); } Some(Ok(Message::Binary(bytes))) => { + self.mark_healthy(); return Some(serde_json::from_slice(&bytes).map_err(Error::from)); } - Some(Ok(Message::Close(_))) | None => return None, - Some(Ok(_)) => continue, - Some(Err(error)) => return Some(Err(Error::WebSocket(error.to_string()))), + Some(Ok(Message::Close(_))) | None => { + if self.reconnect().await { + continue; + } + return None; + } + Some(Ok(_)) => { + // A control frame is not handed to the caller, but it is + // still proof the replacement socket works — and on a quiet + // stream (an inbox with no new items) the server's keepalive + // pings are the *only* traffic. Resetting solely on data + // would let unrelated outages, days apart, accumulate until + // `max_attempts` was spent and the stream ended for good. + self.mark_healthy(); + continue; + } + Some(Err(error)) => { + if self.reconnect().await { + continue; + } + // Prefer why the last re-dial failed over the older error + // that started the run: by now it describes the current + // state of the world, which is what a caller can act on. + return Some(Err(self + .last_dial_error + .take() + .unwrap_or_else(|| Error::WebSocket(error.to_string())))); + } + } + } + } + + /// Record that the socket is carrying traffic: the retry budget for the + /// *next* outage starts over. + /// + /// Deliberately not "reset on open", which is what the TS SDK does: a server + /// that accepts a connection and immediately drops it would then be retried + /// forever. Requiring at least one frame keeps `max_attempts` a real bound + /// on a genuinely broken endpoint. + fn mark_healthy(&mut self) { + self.attempts = 0; + self.last_dial_error = None; + } + + /// Re-dial the stream, reporting whether a fresh one was established. + /// + /// Pauses `interval` *before* each attempt so a server that drops + /// connections instantly cannot be hammered, and bounds each attempt by + /// `connect_timeout` so a stalled handshake cannot hang `recv` forever. + async fn reconnect(&mut self) -> bool { + let policy = self.dialer.reconnect; + if !policy.enabled { + return false; + } + while self.attempts < policy.max_attempts { + self.attempts += 1; + tokio::time::sleep(policy.interval).await; + match tokio::time::timeout(policy.connect_timeout, self.dialer.dial()).await { + Ok(Ok(stream)) => { + self.inner = stream; + self.last_dial_error = None; + return true; + } + Ok(Err(error)) => self.last_dial_error = Some(error), + Err(_) => { + self.last_dial_error = Some(Error::WebSocket(format!( + "reconnect timed out after {:?}", + policy.connect_timeout + ))) + } } } + false } /// Send a JSON message to the server. diff --git a/sdk/rust/tests/websocket.rs b/sdk/rust/tests/websocket.rs index cf9060ce..dc7e0c64 100644 --- a/sdk/rust/tests/websocket.rs +++ b/sdk/rust/tests/websocket.rs @@ -5,6 +5,7 @@ use std::collections::HashMap; use std::sync::mpsc; use std::sync::Arc; +use std::time::Duration; use base64::Engine as _; use ed25519_dalek::{Signature, Verifier, VerifyingKey}; @@ -275,3 +276,286 @@ async fn ws_connect_recv_send_close_round_trip() { conn.close().await.unwrap(); server.await.unwrap(); } + +/// A local server that serves `rounds` successive connections, sending one +/// frame on each before closing it. Returns the bound address and the task; the +/// listener is dropped once every round has been served, so a later dial is +/// refused rather than hanging. +async fn flapping_server(rounds: usize) -> (String, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let task = tokio::spawn(async move { + for n in 1..=rounds { + let (tcp, _) = listener.accept().await.unwrap(); + let mut ws = tokio_tungstenite::accept_async(tcp).await.unwrap(); + ws.send(Message::Text(format!("{{\"type\":\"tick\",\"n\":{n}}}"))) + .await + .unwrap(); + let _ = ws.close(None).await; + } + drop(listener); + }); + (format!("http://{addr}"), task) +} + +fn fast_reconnect(max_attempts: u32) -> tinyplace::ReconnectPolicy { + tinyplace::ReconnectPolicy { + enabled: true, + interval: Duration::from_millis(5), + max_attempts, + connect_timeout: Duration::from_secs(5), + } +} + +#[tokio::test] +async fn ws_reconnect_defaults_match_the_typescript_sdk() { + // Parity with `TinyPlaceWebSocketOptions` in sdk/typescript/src/websocket.ts. + let policy = tinyplace::ReconnectPolicy::default(); + assert!(policy.enabled); + assert_eq!(policy.interval, Duration::from_millis(3_000)); + assert_eq!(policy.max_attempts, 10); + assert!(!tinyplace::ReconnectPolicy::disabled().enabled); +} + +#[tokio::test] +async fn ws_keepalive_traffic_resets_the_retry_budget() { + // A quiet stream — an inbox with no new items — carries nothing but the + // server's keepalive pings. If only data frames reset the budget, outages + // separated by healthy-but-silent connections accumulate until + // `max_attempts` is spent and the stream ends for good. + // + // With `max_attempts: 1`, surviving to the third connection is only + // possible if the pings on connections one and two reset the budget. + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + for round in 1..=3 { + let (tcp, _) = listener.accept().await.unwrap(); + let mut ws = tokio_tungstenite::accept_async(tcp).await.unwrap(); + if round == 3 { + ws.send(Message::Text("{\"n\":3}".to_string())) + .await + .unwrap(); + } else { + ws.send(Message::Ping(Vec::new())).await.unwrap(); + } + let _ = ws.close(None).await; + } + }); + + let client = client(&format!("http://{addr}"), false); + let mut conn = client + .inbox + .stream() + .with_reconnect(fast_reconnect(1)) + .connect() + .await + .unwrap(); + + let frame = tokio::time::timeout(Duration::from_secs(5), conn.recv()) + .await + .expect("recv must not hang") + .expect("the keepalives should have kept the budget alive") + .expect("valid json"); + assert_eq!(frame["n"], 3); + + server.await.unwrap(); +} + +#[tokio::test] +async fn ws_reconnect_bounds_a_stalled_handshake() { + // `connect_async` has no timeout of its own, so a server that accepts the + // TCP connection but never completes the upgrade would hang `recv` forever + // and leave `max_attempts` bounding nothing. + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (tcp, _) = listener.accept().await.unwrap(); + let mut ws = tokio_tungstenite::accept_async(tcp).await.unwrap(); + ws.send(Message::Text("{\"n\":1}".to_string())) + .await + .unwrap(); + let _ = ws.close(None).await; + // Accept, then stall: never answer the upgrade, and hold the socket + // open so the dial cannot fail fast on a closed connection either. + let mut stalled = Vec::new(); + loop { + if let Ok((tcp, _)) = listener.accept().await { + stalled.push(tcp); + } + } + }); + + let client = client(&format!("http://{addr}"), false); + let mut conn = client + .inbox + .stream() + .with_reconnect(tinyplace::ReconnectPolicy { + enabled: true, + interval: Duration::from_millis(5), + max_attempts: 2, + connect_timeout: Duration::from_millis(50), + }) + .connect() + .await + .unwrap(); + + assert_eq!( + conn.recv().await.expect("first frame").expect("valid json")["n"], + 1 + ); + + // Two attempts, each capped at 50ms — comfortably inside this bound. Without + // the per-attempt timeout this hangs until the test harness gives up. + let ended = tokio::time::timeout(Duration::from_secs(5), conn.recv()) + .await + .expect("a stalled handshake must not hang recv"); + assert!(ended.is_none(), "expected the stream to end, got {ended:?}"); + + server.abort(); +} + +#[tokio::test] +async fn ws_recv_reconnects_after_the_server_drops_the_stream() { + let (url, server) = flapping_server(2).await; + let client = client(&url, false); + let mut conn = client + .inbox + .stream() + .with_reconnect(fast_reconnect(5)) + .connect() + .await + .unwrap(); + + // The first frame arrives on the original socket; the server then closes. + let first = conn.recv().await.expect("first frame").expect("valid json"); + assert_eq!(first["n"], 1); + + // The drop is invisible to the caller: recv re-dials and yields the frame + // the *second* connection carries. + let second = conn + .recv() + .await + .expect("second frame") + .expect("valid json"); + assert_eq!(second["n"], 2); + + server.await.unwrap(); +} + +#[tokio::test] +async fn ws_recv_reports_the_drop_when_reconnecting_is_disabled() { + let (url, server) = flapping_server(1).await; + let client = client(&url, false); + let mut conn = client + .inbox + .stream() + .with_reconnect(tinyplace::ReconnectPolicy::disabled()) + .connect() + .await + .unwrap(); + + let first = conn.recv().await.expect("first frame").expect("valid json"); + assert_eq!(first["n"], 1); + // Previous behaviour, preserved as an opt-out: the close ends the stream. + assert!(conn.recv().await.is_none(), "expected the stream to end"); + + server.await.unwrap(); +} + +#[tokio::test] +async fn ws_recv_gives_up_after_max_attempts() { + // One round only: after it, the listener is gone and every re-dial fails. + let (url, server) = flapping_server(1).await; + let client = client(&url, false); + let mut conn = client + .inbox + .stream() + .with_reconnect(fast_reconnect(3)) + .connect() + .await + .unwrap(); + + assert_eq!( + conn.recv().await.expect("first frame").expect("valid json")["n"], + 1 + ); + + // Bounded, not infinite: three failed attempts and it reports the end. + let ended = tokio::time::timeout(Duration::from_secs(5), conn.recv()) + .await + .expect("recv must not hang once attempts are exhausted"); + assert!(ended.is_none(), "expected the stream to end, got {ended:?}"); + + server.await.unwrap(); +} + +// The handshake callback's error type is tungstenite's own `ErrorResponse`. +#[allow(clippy::result_large_err)] +#[tokio::test] +async fn ws_reconnect_signs_a_fresh_upgrade() { + // The upgrade credential is freshness-bound (`v1:::`), so a + // reconnect must re-sign rather than replay the original request — a + // replayed nonce is exactly what the backend rejects with a 401. + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (sender, receiver) = mpsc::channel::>(); + + let server = tokio::spawn(async move { + for n in 1..=2 { + let (tcp, _) = listener.accept().await.unwrap(); + let sender = sender.clone(); + let mut ws = tokio_tungstenite::accept_hdr_async( + tcp, + |request: &Request, response: Response| { + let seen = request + .headers() + .iter() + .map(|(name, value)| { + ( + name.as_str().to_lowercase(), + value.to_str().unwrap_or_default().to_string(), + ) + }) + .collect(); + sender.send(seen).unwrap(); + Ok(response) + }, + ) + .await + .unwrap(); + ws.send(Message::Text(format!("{{\"n\":{n}}}"))) + .await + .unwrap(); + let _ = ws.close(None).await; + } + }); + + let (client, _signer) = client_without_siws(&format!("http://{addr}")); + let mut conn = client + .inbox + .stream() + .with_reconnect(fast_reconnect(5)) + .connect() + .await + .unwrap(); + assert_eq!(conn.recv().await.unwrap().unwrap()["n"], 1); + assert_eq!(conn.recv().await.unwrap().unwrap()["n"], 2); + server.await.unwrap(); + + let first = receiver.recv().unwrap(); + let second = receiver.recv().unwrap(); + let signature = |headers: &HashMap| { + headers + .get("x-tinyplace-signature") + .cloned() + .expect("signature header") + }; + let (first, second) = (signature(&first), signature(&second)); + assert!(first.starts_with("v1:"), "got: {first}"); + assert!(second.starts_with("v1:"), "got: {second}"); + assert_ne!( + first, second, + "the reconnect replayed the original credential instead of re-signing" + ); +}