use std::io;
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};
use rmux_proto::{
encode_frame, ErrorResponse, HandshakeRequest, HasSessionRequest, HasSessionResponse,
ListSessionsRequest, ListSessionsResponse, Request, Response, SdkWaitId, SessionName,
};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};
use tokio::task::JoinHandle;
use tokio::time::{timeout, Duration};
use super::{allocate_bounded_atomic_id, mix_sdk_wait_owner_id, DropGuard, TransportClient};
use crate::RmuxError;
fn alpha() -> SessionName {
SessionName::new("alpha").expect("valid session")
}
fn has_session_request() -> Request {
Request::HasSession(HasSessionRequest { target: alpha() })
}
fn list_sessions_request() -> Request {
Request::ListSessions(ListSessionsRequest {
format: None,
filter: None,
sort_order: None,
reversed: false,
})
}
fn handshake_request() -> Request {
Request::Handshake(HandshakeRequest::requiring(["capability.future"]))
}
fn has_session_response(exists: bool) -> Response {
Response::HasSession(HasSessionResponse { exists })
}
fn list_sessions_response(stdout: &[u8]) -> Response {
Response::ListSessions(ListSessionsResponse {
output: rmux_proto::CommandOutput::from_stdout(stdout),
})
}
fn session_not_found_response() -> Response {
Response::Error(ErrorResponse {
error: rmux_proto::RmuxError::SessionNotFound("alpha".to_owned()),
})
}
fn unsupported_capability_response() -> Response {
Response::Error(ErrorResponse {
error: rmux_proto::RmuxError::UnsupportedCapability {
feature: "capability.future".to_owned(),
supported: vec![rmux_proto::CAPABILITY_HANDSHAKE.to_owned()],
},
})
}
fn unsupported_wire_version_response() -> Response {
Response::Error(ErrorResponse {
error: rmux_proto::RmuxError::UnsupportedWireVersion {
got: rmux_proto::RMUX_WIRE_VERSION + 1,
minimum: rmux_proto::RMUX_WIRE_VERSION,
maximum: rmux_proto::RMUX_WIRE_VERSION,
},
})
}
async fn read_request(stream: &mut tokio::io::DuplexStream) -> Request {
let mut decoder = rmux_proto::FrameDecoder::new();
let mut buffer = [0; 256];
loop {
if let Some(request) = decoder
.next_frame::<Request>()
.expect("request frame decodes")
{
return request;
}
let read = stream.read(&mut buffer).await.expect("read request bytes");
assert_ne!(read, 0, "client closed before request arrived");
decoder.push_bytes(&buffer[..read]);
}
}
async fn write_response(stream: &mut tokio::io::DuplexStream, response: &Response) {
let frame = encode_frame(response).expect("response encodes");
stream.write_all(&frame).await.expect("write response");
stream.flush().await.expect("flush response");
}
async fn write_unknown_response_tag(stream: &mut tokio::io::DuplexStream) {
let payload = 255_u32.to_le_bytes();
let mut frame = vec![
rmux_proto::RMUX_FRAME_MAGIC,
rmux_proto::RMUX_WIRE_VERSION as u8,
];
frame.extend_from_slice(&(payload.len() as u32).to_le_bytes());
frame.extend_from_slice(&payload);
stream.write_all(&frame).await.expect("write bad frame");
stream.flush().await.expect("flush bad frame");
}
fn assert_transport_kind(result: crate::Result<Response>, expected: io::ErrorKind) {
match result.expect_err("request must fail") {
RmuxError::Transport { source, .. } => assert_eq!(source.kind(), expected),
error => panic!("expected transport error, got {error:?}"),
}
}
fn assert_transport_message(
result: crate::Result<Response>,
expected: io::ErrorKind,
expected_message: &str,
) {
match result.expect_err("request must fail") {
RmuxError::Transport { source, .. } => {
assert_eq!(source.kind(), expected);
assert!(
source.to_string().contains(expected_message),
"transport source `{source}` must contain `{expected_message}`"
);
}
error => panic!("expected transport error, got {error:?}"),
}
}
fn assert_unsupported_feature(result: crate::Result<Response>, expected: &str) {
match result.expect_err("request must fail") {
RmuxError::Unsupported { feature, .. } => assert_eq!(feature, expected),
error => panic!("expected unsupported error, got {error:?}"),
}
}
#[tokio::test]
async fn sdk_wait_owner_ids_are_distinct_and_wait_ids_are_per_transport() {
let (first_stream, _first_peer) = tokio::io::duplex(64);
let (second_stream, _second_peer) = tokio::io::duplex(64);
let first = TransportClient::spawn(first_stream);
let second = TransportClient::spawn(second_stream);
assert_ne!(first.sdk_wait_owner_id(), second.sdk_wait_owner_id());
assert_eq!(first.allocate_sdk_wait_id(), SdkWaitId::new(1));
assert_eq!(first.allocate_sdk_wait_id(), SdkWaitId::new(2));
assert_eq!(second.allocate_sdk_wait_id(), SdkWaitId::new(1));
}
#[test]
fn sdk_wait_owner_ids_mix_process_seed_with_local_counter() {
assert_ne!(
mix_sdk_wait_owner_id(0x1111_2222_3333_4444, 1),
mix_sdk_wait_owner_id(0x5555_6666_7777_8888, 1)
);
assert_ne!(
mix_sdk_wait_owner_id(0x1111_2222_3333_4444, 1),
mix_sdk_wait_owner_id(0x1111_2222_3333_4444, 2)
);
assert_ne!(mix_sdk_wait_owner_id(0, 0), 0);
}
#[test]
fn bounded_sdk_wait_id_allocator_panics_instead_of_wrapping() {
let counter = AtomicU64::new(2);
assert_eq!(allocate_bounded_atomic_id(&counter, 2, "exhausted"), 2);
assert_eq!(counter.load(Ordering::Relaxed), 3);
let result = catch_unwind(AssertUnwindSafe(|| {
allocate_bounded_atomic_id(&counter, 2, "exhausted");
}));
assert!(result.is_err());
assert_eq!(counter.load(Ordering::Relaxed), 3);
}
fn spawn_request(
client: &TransportClient,
request: Request,
) -> JoinHandle<crate::Result<Response>> {
let client = client.clone();
tokio::spawn(async move { client.request(request).await })
}
async fn join_request(handle: JoinHandle<crate::Result<Response>>) -> crate::Result<Response> {
handle.await.expect("request task must not panic")
}
#[tokio::test]
async fn actor_correlates_bare_responses_in_fifo_request_order() {
let (client_stream, mut server_stream) = tokio::io::duplex(4096);
let client = TransportClient::spawn(client_stream);
let first = spawn_request(&client, has_session_request());
assert_eq!(
read_request(&mut server_stream).await,
has_session_request()
);
let second = spawn_request(&client, list_sessions_request());
assert_eq!(
read_request(&mut server_stream).await,
list_sessions_request()
);
write_response(&mut server_stream, &has_session_response(true)).await;
write_response(&mut server_stream, &list_sessions_response(b"alpha\n")).await;
let (first, second) = tokio::join!(join_request(first), join_request(second));
assert_eq!(first.expect("first response"), has_session_response(true));
assert_eq!(
second.expect("second response"),
list_sessions_response(b"alpha\n")
);
}
#[tokio::test]
async fn actor_rejects_out_of_order_response_kinds_and_closes_transport() {
let (client_stream, mut server_stream) = tokio::io::duplex(4096);
let client = TransportClient::spawn(client_stream);
let first = spawn_request(&client, has_session_request());
assert_eq!(
read_request(&mut server_stream).await,
has_session_request()
);
let second = spawn_request(&client, list_sessions_request());
assert_eq!(
read_request(&mut server_stream).await,
list_sessions_request()
);
write_response(&mut server_stream, &list_sessions_response(b"alpha\n")).await;
let (first, second) = tokio::join!(join_request(first), join_request(second));
assert_transport_message(
first,
io::ErrorKind::InvalidData,
"sent `list-sessions` response for pending `has-session` request",
);
assert_transport_kind(second, io::ErrorKind::InvalidData);
assert_transport_kind(
client.request(has_session_request()).await,
io::ErrorKind::InvalidData,
);
}
#[tokio::test]
async fn error_response_completes_current_fifo_slot_as_protocol_error() {
let (client_stream, mut server_stream) = tokio::io::duplex(4096);
let client = TransportClient::spawn(client_stream);
let request = spawn_request(&client, has_session_request());
assert_eq!(
read_request(&mut server_stream).await,
has_session_request()
);
write_response(&mut server_stream, &session_not_found_response()).await;
match join_request(request).await.expect_err("request must fail") {
RmuxError::Protocol { source } => {
assert_eq!(
source,
rmux_proto::RmuxError::SessionNotFound("alpha".to_owned())
);
}
error => panic!("expected protocol error, got {error:?}"),
}
}
#[tokio::test]
async fn unsupported_capability_response_maps_to_stable_sdk_feature() {
let (client_stream, mut server_stream) = tokio::io::duplex(4096);
let client = TransportClient::spawn(client_stream);
let request = spawn_request(&client, handshake_request());
assert_eq!(read_request(&mut server_stream).await, handshake_request());
write_response(&mut server_stream, &unsupported_capability_response()).await;
match join_request(request).await.expect_err("request must fail") {
RmuxError::Unsupported { feature, hint } => {
assert_eq!(feature, "capability.future");
assert!(hint.contains(rmux_proto::CAPABILITY_HANDSHAKE));
}
error => panic!("expected unsupported error, got {error:?}"),
}
}
#[tokio::test]
async fn unsupported_wire_version_response_maps_to_stable_sdk_feature() {
let (client_stream, mut server_stream) = tokio::io::duplex(4096);
let client = TransportClient::spawn(client_stream);
let request = spawn_request(&client, handshake_request());
assert_eq!(read_request(&mut server_stream).await, handshake_request());
write_response(&mut server_stream, &unsupported_wire_version_response()).await;
match join_request(request).await.expect_err("request must fail") {
RmuxError::Unsupported { feature, .. } => {
assert_eq!(feature, crate::FEATURE_PROTOCOL_WIRE_VERSION);
}
error => panic!("expected unsupported error, got {error:?}"),
}
}
#[tokio::test]
async fn handshake_response_decode_mismatch_maps_to_stable_sdk_feature() {
let (client_stream, mut server_stream) = tokio::io::duplex(4096);
let client = TransportClient::spawn(client_stream);
let request = spawn_request(&client, handshake_request());
assert_eq!(read_request(&mut server_stream).await, handshake_request());
write_unknown_response_tag(&mut server_stream).await;
assert_unsupported_feature(
join_request(request).await,
crate::FEATURE_PROTOCOL_CAPABILITIES,
);
}
#[tokio::test]
async fn unsolicited_response_without_pending_request_closes_transport() {
let (client_stream, mut server_stream) = tokio::io::duplex(4096);
let client = TransportClient::spawn(client_stream);
write_response(&mut server_stream, &has_session_response(true)).await;
timeout(Duration::from_secs(1), async {
while client.state.terminal_failure().is_none() {
tokio::task::yield_now().await;
}
})
.await
.expect("actor must reject unsolicited response before the next request");
assert_transport_message(
client.request(has_session_request()).await,
io::ErrorKind::InvalidData,
"sent unsolicited `has-session` response",
);
}
#[tokio::test]
async fn transport_shutdown_waits_for_peer_close() {
let (client_stream, mut server_stream) = tokio::io::duplex(4096);
let client = TransportClient::spawn(client_stream);
let client_for_shutdown = client.clone();
let mut shutdown = tokio::spawn(async move { client_for_shutdown.shutdown().await });
let mut buffer = [0_u8; 1];
assert_eq!(
server_stream
.read(&mut buffer)
.await
.expect("read client eof"),
0
);
assert!(
timeout(Duration::from_millis(50), &mut shutdown)
.await
.is_err(),
"shutdown must wait for the peer read side to close"
);
drop(server_stream);
shutdown
.await
.expect("shutdown task")
.expect("shutdown succeeds after peer close");
}
#[tokio::test]
async fn transport_shutdown_treats_prior_peer_eof_as_clean_close() {
let (client_stream, mut server_stream) = tokio::io::duplex(4096);
let client = TransportClient::spawn(client_stream);
let request = spawn_request(&client, has_session_request());
assert_eq!(
read_request(&mut server_stream).await,
has_session_request()
);
write_response(&mut server_stream, &has_session_response(true)).await;
drop(server_stream);
assert_eq!(
join_request(request).await.expect("request response"),
has_session_response(true)
);
timeout(Duration::from_secs(1), async {
while !client
.state
.terminal_failure()
.is_some_and(|failure| failure.is_eof())
{
tokio::task::yield_now().await;
}
})
.await
.expect("transport observes peer close");
client
.shutdown()
.await
.expect("already closed peer is a clean shutdown");
}
#[tokio::test]
async fn actor_wakes_every_pending_caller_on_eof() {
let (client_stream, mut server_stream) = tokio::io::duplex(4096);
let client = TransportClient::spawn(client_stream);
let first = spawn_request(&client, has_session_request());
assert_eq!(
read_request(&mut server_stream).await,
has_session_request()
);
let second = spawn_request(&client, list_sessions_request());
assert_eq!(
read_request(&mut server_stream).await,
list_sessions_request()
);
drop(server_stream);
let (first, second) = tokio::join!(join_request(first), join_request(second));
assert_transport_kind(first, io::ErrorKind::UnexpectedEof);
assert_transport_kind(second, io::ErrorKind::UnexpectedEof);
assert_transport_kind(
client.request(has_session_request()).await,
io::ErrorKind::UnexpectedEof,
);
}
#[tokio::test]
async fn actor_wakes_every_pending_caller_on_bad_frame() {
let (client_stream, mut server_stream) = tokio::io::duplex(4096);
let client = TransportClient::spawn(client_stream);
let first = spawn_request(&client, has_session_request());
assert_eq!(
read_request(&mut server_stream).await,
has_session_request()
);
let second = spawn_request(&client, list_sessions_request());
assert_eq!(
read_request(&mut server_stream).await,
list_sessions_request()
);
server_stream
.write_all(&[0])
.await
.expect("write invalid frame byte");
let (first, second) = tokio::join!(join_request(first), join_request(second));
assert_transport_kind(first, io::ErrorKind::InvalidData);
assert_transport_kind(second, io::ErrorKind::InvalidData);
}
#[tokio::test]
async fn actor_wakes_every_pending_caller_on_read_error() {
let client = TransportClient::spawn(ScriptedIo::read_error_after_writes(2));
let first = client.request(has_session_request());
let second = client.request(list_sessions_request());
let (first, second) = tokio::join!(first, second);
assert_transport_kind(first, io::ErrorKind::ConnectionReset);
assert_transport_kind(second, io::ErrorKind::ConnectionReset);
assert_transport_kind(
client.request(has_session_request()).await,
io::ErrorKind::ConnectionReset,
);
}
#[tokio::test]
async fn actor_wakes_every_pending_caller_on_write_error() {
let client = TransportClient::spawn(ScriptedIo::write_error_on_call(2));
let first = client.request(has_session_request());
let second = client.request(list_sessions_request());
let (first, second) = tokio::join!(first, second);
assert_transport_kind(first, io::ErrorKind::BrokenPipe);
assert_transport_kind(second, io::ErrorKind::BrokenPipe);
assert_transport_kind(
client.request(has_session_request()).await,
io::ErrorKind::BrokenPipe,
);
}
#[tokio::test]
async fn terminal_read_error_before_request_write_is_reported_explicitly() {
let client = TransportClient::spawn(ScriptedIo::read_error_after_writes(0));
assert_transport_kind(
client.request(has_session_request()).await,
io::ErrorKind::ConnectionReset,
);
assert_transport_kind(
client.request(list_sessions_request()).await,
io::ErrorKind::ConnectionReset,
);
}
#[tokio::test]
async fn drop_guard_uses_nonblocking_best_effort_actor_send() {
let (client_stream, mut server_stream) = tokio::io::duplex(4096);
let client = TransportClient::spawn(client_stream);
drop(DropGuard::best_effort(
client.clone(),
has_session_request(),
));
assert_eq!(
read_request(&mut server_stream).await,
has_session_request()
);
write_response(&mut server_stream, &has_session_response(true)).await;
let follow_up = spawn_request(&client, list_sessions_request());
assert_eq!(
read_request(&mut server_stream).await,
list_sessions_request()
);
write_response(&mut server_stream, &list_sessions_response(b"alpha\n")).await;
assert_eq!(
join_request(follow_up).await.expect("follow-up response"),
list_sessions_response(b"alpha\n")
);
}
#[derive(Clone, Copy)]
enum Script {
ReadErrorAfterWrites { writes: usize },
WriteErrorOnCall { call: usize },
}
struct ScriptedIo {
script: Script,
state: Arc<Mutex<ScriptedIoState>>,
}
#[derive(Default)]
struct ScriptedIoState {
write_calls: usize,
read_waker: Option<Waker>,
}
impl ScriptedIo {
fn read_error_after_writes(writes: usize) -> Self {
Self {
script: Script::ReadErrorAfterWrites { writes },
state: Arc::new(Mutex::new(ScriptedIoState::default())),
}
}
fn write_error_on_call(call: usize) -> Self {
Self {
script: Script::WriteErrorOnCall { call },
state: Arc::new(Mutex::new(ScriptedIoState::default())),
}
}
}
impl AsyncRead for ScriptedIo {
fn poll_read(
self: Pin<&mut Self>,
context: &mut Context<'_>,
_buffer: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let mut state = self.state.lock().expect("scripted state lock");
match self.script {
Script::ReadErrorAfterWrites { writes } if state.write_calls >= writes => {
Poll::Ready(Err(io::Error::new(
io::ErrorKind::ConnectionReset,
"scripted read error",
)))
}
_ => {
state.read_waker = Some(context.waker().clone());
Poll::Pending
}
}
}
}
impl AsyncWrite for ScriptedIo {
fn poll_write(
self: Pin<&mut Self>,
_context: &mut Context<'_>,
buffer: &[u8],
) -> Poll<io::Result<usize>> {
let mut state = self.state.lock().expect("scripted state lock");
state.write_calls += 1;
if matches!(self.script, Script::WriteErrorOnCall { call } if state.write_calls == call) {
return Poll::Ready(Err(io::Error::new(
io::ErrorKind::BrokenPipe,
"scripted write error",
)));
}
if let Some(waker) = state.read_waker.take() {
waker.wake();
}
Poll::Ready(Ok(buffer.len()))
}
fn poll_flush(self: Pin<&mut Self>, _context: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_shutdown(self: Pin<&mut Self>, _context: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
}