mod common;
use std::time::Duration;
use common::TestCall;
use futures_util::{SinkExt, StreamExt};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tokio_tungstenite::tungstenite::Message;
use unb::{handler, Handler, HandlerError, Reply, Request};
use unb_core::{Envelope, Kind, PROTOCOL_VERSION};
use unb_server::{Endpoint, EndpointSet, Node, TransportKind};
use unb_server::{HostConfig, TcpTransport};
#[derive(Deserialize, JsonSchema)]
struct Probe {}
#[derive(Deserialize, Serialize, JsonSchema)]
#[serde(transparent)]
struct EchoPayload(Value);
#[handler]
async fn weather(_request: Request<Probe>) -> Result<Reply<Value>, HandlerError> {
Ok(Reply::new(json!({ "temp_c": 21 })))
}
#[handler]
async fn a_echo(request: Request<EchoPayload>) -> Result<Reply<Value>, HandlerError> {
Ok(Reply::new(json!({ "from": "a", "got": request.payload() })))
}
#[handler]
async fn b_echo(request: Request<EchoPayload>) -> Result<Reply<Value>, HandlerError> {
Ok(Reply::new(json!({ "from": "b", "got": request.payload() })))
}
#[derive(Clone, Copy)]
enum ScriptMode {
StallBeforeWelcome,
StallAfterIdentity,
}
async fn scripted_ws_server(mode: ScriptMode) -> String {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
while let Ok((stream, _)) = listener.accept().await {
tokio::spawn(async move {
let Ok(socket) = tokio_tungstenite::accept_async(stream).await else {
return;
};
let (mut sink, mut source) = socket.split();
let mut next_frame = 0u64;
let mut frame = |kind: Kind, payload: &Value| {
next_frame += 1;
Envelope {
v: PROTOCOL_VERSION,
id: format!("g{next_frame}"),
target: String::new(),
subject: String::new(),
kind,
corr: None,
seq: None,
hops: None,
body_token: None,
payload: Envelope::encode_payload(payload),
path: Vec::new(),
headers: Default::default(),
}
.encode()
};
while let Some(Ok(message)) = source.next().await {
let Message::Binary(bytes) = message else {
continue;
};
let Ok(envelope) = Envelope::decode(bytes) else {
continue;
};
match (mode, envelope.kind) {
(ScriptMode::StallBeforeWelcome, Kind::Hello) => {}
(ScriptMode::StallAfterIdentity, Kind::Hello) => {
let welcome = frame(Kind::Welcome, &json!({ "version": 1 }));
let _ = sink.send(Message::Binary(welcome.to_vec().into())).await;
}
(ScriptMode::StallAfterIdentity, Kind::Identify) => {
let identity = frame(
Kind::Identify,
&json!({
"node_id": "staller",
"instance_id": "staller-1",
"epoch": 1
}),
);
let _ = sink.send(Message::Binary(identity.to_vec().into())).await;
let accepted = frame(Kind::IdentityAccepted, &Value::Null);
let _ = sink.send(Message::Binary(accepted.to_vec().into())).await;
}
_ => {}
}
}
});
}
});
format!("ws://{addr}")
}
fn ws_endpoint(address: String) -> Endpoint {
Endpoint {
kind: TransportKind::WebSocket,
address,
cert_hash: None,
}
}
async fn hosted_weather_node() -> (std::sync::Arc<Node>, String, unb_server::Hosting) {
let node = Node::builder("weather-1")
.service(weather)
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let hosting = HostConfig::tcp(([127, 0, 0, 1], 0), TcpTransport::plain())
.start(&node)
.await
.unwrap();
let url = format!("ws://{}", hosting.websocket_addr().unwrap());
(node, url, hosting)
}
fn fast_dialer(name: &str) -> std::sync::Arc<Node> {
Node::builder(name)
.connect_timeout(Duration::from_millis(500))
.insecure_accept_declared_peer_identities()
.build()
.unwrap()
}
fn unix_endpoint(path: &std::path::Path) -> Endpoint {
Endpoint {
kind: TransportKind::Unix,
address: path.to_string_lossy().into_owned(),
cert_hash: None,
}
}
fn unix_serve(node: &std::sync::Arc<Node>, dir: &tempfile::TempDir) -> std::path::PathBuf {
let path = dir.path().join("owner.sock");
let listener = unb_transport::unix::UnixListener::bind(&path).unwrap();
let node = node.clone();
tokio::spawn(async move {
loop {
let Ok(pipe) = listener.accept().await else {
return;
};
let node = node.clone();
let bodies = listener.streams();
tokio::spawn(async move {
node.serve_transport(unb_runtime::Pipe::piped_with_streams(pipe, false, bodies))
.await;
});
}
});
path
}
#[tokio::test(flavor = "multi_thread")]
async fn a_unix_endpoint_outranks_websocket_for_the_same_peer() {
let dir = tempfile::tempdir().unwrap();
let node = Node::builder("weather-1")
.service(weather)
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let socket = unix_serve(&node, &dir);
let hosting = HostConfig::tcp(([127, 0, 0, 1], 0), TcpTransport::plain())
.start(&node)
.await
.unwrap();
let ws_url = format!("ws://{}", hosting.websocket_addr().unwrap());
let dialer = fast_dialer("caller-1");
let set = EndpointSet::from(vec![ws_endpoint(ws_url), unix_endpoint(&socket)]);
dialer.connect(set).await.unwrap();
assert!(dialer.reachable_names().contains(&"weather-1".to_string()));
}
#[tokio::test(flavor = "multi_thread")]
async fn a_dead_unix_endpoint_falls_back_to_websocket() {
let dir = tempfile::tempdir().unwrap();
let (_owner, good_url, _hosting) = hosted_weather_node().await;
let dialer = fast_dialer("caller-1");
let set = EndpointSet::from(vec![
unix_endpoint(&dir.path().join("absent.sock")),
ws_endpoint(good_url),
]);
dialer.connect(set).await.unwrap();
assert!(dialer.reachable_names().contains(&"weather-1".to_string()));
}
#[tokio::test(flavor = "multi_thread")]
async fn a_dead_preferred_endpoint_falls_back_to_the_working_one() {
let (_owner, good_url, _hosting) = hosted_weather_node().await;
let dialer = fast_dialer("caller-1");
let set = EndpointSet::from(vec![
ws_endpoint("ws://127.0.0.1:9".into()),
ws_endpoint(good_url),
]);
dialer.connect(set).await.unwrap();
assert!(
dialer.reachable_names().contains(&"weather-1".to_string()),
"routes are learned as part of readiness"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_handshake_stall_times_out_and_falls_back() {
let stalling = scripted_ws_server(ScriptMode::StallBeforeWelcome).await;
let (_owner, good_url, _hosting) = hosted_weather_node().await;
let dialer = fast_dialer("caller-1");
let set = EndpointSet::from(vec![ws_endpoint(stalling), ws_endpoint(good_url)]);
tokio::time::timeout(Duration::from_secs(7), dialer.connect(set))
.await
.unwrap()
.unwrap();
assert!(dialer.reachable_names().contains(&"weather-1".to_string()));
}
#[tokio::test(flavor = "multi_thread")]
async fn a_route_sync_stall_times_out_and_falls_back() {
let stalling = scripted_ws_server(ScriptMode::StallAfterIdentity).await;
let (_owner, good_url, _hosting) = hosted_weather_node().await;
let dialer = fast_dialer("caller-1");
let set = EndpointSet::from(vec![ws_endpoint(stalling), ws_endpoint(good_url)]);
dialer.connect(set).await.unwrap();
assert!(dialer.reachable_names().contains(&"weather-1".to_string()));
}
#[tokio::test(flavor = "multi_thread")]
async fn endpoint_attempt_timeout_does_not_decide_protocol_establishment() {
let stalling = scripted_ws_server(ScriptMode::StallBeforeWelcome).await;
let dialer = Node::builder("caller-1")
.connect_timeout(Duration::from_millis(50))
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let connecting = tokio::spawn({
let dialer = dialer.clone();
async move {
dialer
.connect(EndpointSet::from(ws_endpoint(stalling)))
.await
}
});
tokio::time::sleep(Duration::from_millis(150)).await;
assert!(
!connecting.is_finished(),
"endpoint dialing policy must not decide protocol establishment timeout"
);
connecting.abort();
}
#[tokio::test(flavor = "multi_thread")]
async fn every_candidate_failing_returns_the_final_error_and_leaves_no_state() {
let stalling = scripted_ws_server(ScriptMode::StallBeforeWelcome).await;
let dialer = fast_dialer("caller-1");
let set = EndpointSet::from(vec![
ws_endpoint("ws://127.0.0.1:9".into()),
ws_endpoint(stalling),
]);
let result = dialer.connect(set).await;
assert!(result.is_err());
assert_eq!(
dialer.reachable_names(),
["caller-1"],
"no failed candidate leaves a remote route"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn simultaneous_cross_connect_converges_and_serves_both_directions() {
let a = Node::builder("node-a")
.service(a_echo.at_subject("a.echo"))
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let b = Node::builder("node-b")
.service(b_echo.at_subject("b.echo"))
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let a_hosting = HostConfig::tcp(([127, 0, 0, 1], 0), TcpTransport::plain())
.start(&a)
.await
.unwrap();
let b_hosting = HostConfig::tcp(([127, 0, 0, 1], 0), TcpTransport::plain())
.start(&b)
.await
.unwrap();
let a_url = format!("ws://{}", a_hosting.websocket_addr().unwrap());
let b_url = format!("ws://{}", b_hosting.websocket_addr().unwrap());
let (from_a, from_b) = tokio::join!(
a.connect(EndpointSet::from(ws_endpoint(b_url))),
b.connect(EndpointSet::from(ws_endpoint(a_url)))
);
let from_a = from_a.unwrap();
let from_b = from_b.unwrap();
assert_eq!(from_a.peer(), "node-b");
assert_eq!(from_b.peer(), "node-a");
assert_eq!(from_a.status(), unb_server::ConnectionStatus::Connected);
assert_eq!(from_b.status(), unb_server::ConnectionStatus::Connected);
common::wait_until("a learns node-b", || {
a.reachable_names().contains(&"node-b".to_string())
})
.await;
common::wait_until("b learns node-a", || {
b.reachable_names().contains(&"node-a".to_string())
})
.await;
tokio::time::sleep(Duration::from_millis(200)).await;
common::wait_until("a still reaches node-b after loser cleanup", || {
a.reachable_names().contains(&"node-b".to_string())
})
.await;
common::wait_until("b still reaches node-a after loser cleanup", || {
b.reachable_names().contains(&"node-a".to_string())
})
.await;
let from_a = a
.request("/node-b/b.echo", json!({ "direction": "a-to-b" }))
.await
.unwrap();
assert_eq!(
from_a,
json!({ "from": "b", "got": { "direction": "a-to-b" } })
);
let from_b = b
.request("/node-a/a.echo", json!({ "direction": "b-to-a" }))
.await
.unwrap();
assert_eq!(
from_b,
json!({ "from": "a", "got": { "direction": "b-to-a" } })
);
}