mod common;
use common::TestCall;
use std::time::Duration;
use bytes::Bytes;
use common::start;
use futures_util::{SinkExt, StreamExt};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use unb::{handler, Handler, HandlerError, Reply, Request};
use unb_client::{dial_transport, pair};
use unb_core::Kind;
use unb_runtime::{ClientError, Wire};
use unb_server::Node;
#[derive(Deserialize, Serialize, JsonSchema)]
#[serde(transparent)]
struct EchoPayload(Value);
#[derive(Deserialize, JsonSchema)]
struct Probe {}
#[handler]
async fn echo(request: Request<EchoPayload>) -> Result<Reply<EchoPayload>, HandlerError> {
Ok(Reply::new(request.into_payload()))
}
#[handler]
async fn boom(_request: Request<Probe>) -> Result<Reply<Value>, HandlerError> {
panic!("handler exploded")
}
#[handler]
async fn ok_fine(_request: Request<Probe>) -> Result<Reply<Value>, HandlerError> {
Ok(Reply::new(json!({ "fine": true })))
}
#[handler]
async fn block(_request: Request<Probe>) -> Result<Reply<Value>, HandlerError> {
std::future::pending::<()>().await;
Ok(Reply::new(json!({})))
}
#[tokio::test(flavor = "multi_thread")]
async fn a_malformed_binary_frame_closes_the_session() {
let node = Node::builder("echo-1")
.service(echo)
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let url = start(node).await;
let (mut socket, _) = tokio_tungstenite::connect_async(&url).await.unwrap();
socket
.send(tokio_tungstenite::tungstenite::Message::Binary(
Bytes::from_static(b"\xff\xff\xff\xffgarbage"),
))
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(5), async {
loop {
match socket.next().await {
None
| Some(Err(_))
| Some(Ok(tokio_tungstenite::tungstenite::Message::Close(_))) => return,
Some(Ok(_)) => {}
}
}
})
.await
.expect("a malformed frame must close the session");
}
#[tokio::test(flavor = "multi_thread")]
async fn dropping_the_last_node_handle_shuts_down_its_sessions() {
let node = Node::builder("ephemeral-1")
.service(echo)
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let (client_side, server) = pair();
node.serve_transport(server).await;
let client = Wire::open(client_side);
drop(node);
let ended = tokio::time::timeout(Duration::from_secs(5), async {
client.closed().await;
})
.await;
assert!(
ended.is_ok(),
"dropping the last Node handle must stop its sessions via DropGuard"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_panicking_handler_leaves_the_node_serving() {
let node = Node::builder("sturdy-1")
.service(boom)
.service(ok_fine.at_subject("ok"))
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let url = start(node).await;
let wire = Wire::open(dial_transport(&url).await.unwrap());
let _boom = common::stream(&wire, "/sturdy-1/boom", Kind::Request, json!({})).await;
let mut ok = common::stream(&wire, "/sturdy-1/ok", Kind::Request, json!({})).await;
let answered = tokio::time::timeout(Duration::from_secs(5), ok.next())
.await
.expect("the node must keep serving after a handler panic")
.unwrap()
.unwrap();
assert_eq!(
answered.payload_json()["fine"],
true,
"no lock poisoning, no dead node"
);
}
#[tokio::test(start_paused = true)]
async fn a_silent_unclassified_session_is_closed_at_the_establishment_deadline() {
let node = Node::builder("deadline-1")
.service(echo)
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let (client_side, server) = pair();
node.serve_transport(server).await;
let client = Wire::open(client_side);
let closed = tokio::time::timeout(Duration::from_secs(30), client.closed()).await;
assert!(
closed.is_ok(),
"a session that never classifies must be closed at the establishment deadline"
);
}
#[tokio::test(start_paused = true)]
async fn a_half_established_peer_is_closed_at_the_establishment_deadline() {
let node = Node::builder("deadline-2")
.service(echo)
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let (client_side, server) = pair();
node.serve_transport(server).await;
let ghost = Wire::open(client_side);
let identity = unb_core::NodeIdentity {
node_id: "ghost-1".into(),
instance_id: "ghost-1-1".into(),
epoch: 1,
proof: json!(null),
};
ghost
.control(
Kind::Identify,
unb_core::Envelope::encode_payload(&serde_json::to_value(&identity).unwrap()),
)
.await
.unwrap();
let closed = tokio::time::timeout(Duration::from_secs(30), ghost.closed()).await;
assert!(
closed.is_ok(),
"a peer stuck mid-establishment must be closed at the deadline"
);
assert!(
!node.reachable_names().iter().any(|name| name == "ghost-1"),
"a half-established peer must contribute no route state"
);
}
#[tokio::test(start_paused = true)]
async fn a_classified_client_session_survives_past_the_establishment_deadline() {
let node = Node::builder("deadline-3")
.service(echo)
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let (client_side, server) = pair();
node.serve_transport(server).await;
let client = Wire::open(client_side);
let mut call = common::stream(
&client,
"/deadline-3/echo",
Kind::Request,
json!({ "n": 1 }),
)
.await;
let first = tokio::time::timeout(Duration::from_secs(5), call.next())
.await
.expect("the first request is answered")
.unwrap()
.unwrap();
assert_eq!(first.payload_json()["n"], 1);
tokio::time::sleep(Duration::from_secs(8)).await;
let mut call = common::stream(
&client,
"/deadline-3/echo",
Kind::Request,
json!({ "n": 2 }),
)
.await;
let second = tokio::time::timeout(Duration::from_secs(5), call.next())
.await
.expect("the second request is answered")
.unwrap()
.unwrap();
assert_eq!(second.payload_json()["n"], 2);
}
#[tokio::test(start_paused = true)]
async fn a_ready_peer_link_survives_past_the_establishment_deadline() {
let a = Node::builder("deadline-a")
.service(echo)
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let b = Node::builder("deadline-b")
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
b.link(&a).await.unwrap();
tokio::time::sleep(Duration::from_secs(8)).await;
let value = b
.request("/deadline-a/echo", json!({ "n": 3 }))
.await
.unwrap();
assert_eq!(
value["n"], 3,
"a ready peer link must survive past the establishment deadline"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_request_flood_past_the_activation_bound_is_refused_with_busy() {
let node = Node::builder("bounded-1")
.max_activations(2)
.service(block)
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let url = start(node).await;
let wire = Wire::open(dial_transport(&url).await.unwrap());
let mut calls = futures_util::stream::FuturesUnordered::new();
for _ in 0..3 {
let mut call = common::stream(&wire, "/bounded-1/block", Kind::Request, json!({})).await;
calls.push(async move { call.next().await });
}
let busy = tokio::time::timeout(Duration::from_secs(5), async {
while let Some(result) = calls.next().await {
if matches!(
result,
Err(ClientError::Protocol {
code: unb_core::ErrorCode::Busy,
..
})
) {
return true;
}
}
false
})
.await
.expect("the over-capacity request is answered promptly");
assert!(
busy,
"a request past the activation bound is refused with BUSY, never buffered"
);
}