use alloc::format;
use alloc::string::ToString;
use alloc::vec;
use alloc::vec::Vec;
use core::time::Duration;
use std::net::TcpListener;
use std::sync::mpsc;
use liminal::protocol::{Frame, ProtocolVersion, SchemaId, encode, encoded_len};
use tungstenite::Message;
use super::WebSocketSubscriptionStream;
use crate::SdkError;
use crate::remote::mailbox::MailboxLimits;
const GUARD: Duration = Duration::from_secs(60);
fn encoded(frame: &Frame) -> Result<Vec<u8>, SdkError> {
let protocol = |error: liminal::protocol::ProtocolError| SdkError::Protocol {
description: error.to_string(),
};
let mut bytes = vec![0_u8; encoded_len(frame).map_err(protocol)?];
let written = encode(frame, &mut bytes).map_err(protocol)?;
bytes.truncate(written);
Ok(bytes)
}
fn server_error(error: &impl core::fmt::Display) -> SdkError {
SdkError::Connection {
description: format!("fake websocket server: {error}"),
}
}
fn malformed(frame_type: u8) -> Vec<u8> {
let mut bytes = vec![frame_type, 0, 0, 0, 0, 1];
bytes.extend_from_slice(&3_u32.to_be_bytes());
bytes.extend_from_slice(&[0xEE; 3]);
bytes
}
fn serve(
listener: TcpListener,
gate: Option<mpsc::Receiver<()>>,
after: Vec<Vec<u8>>,
) -> std::thread::JoinHandle<Result<(), SdkError>> {
std::thread::spawn(move || -> Result<(), SdkError> {
let (stream, _) = listener.accept().map_err(|error| server_error(&error))?;
stream
.set_read_timeout(Some(GUARD))
.map_err(|error| server_error(&error))?;
let mut socket = tungstenite::accept(stream).map_err(|error| server_error(&error))?;
let replies = [
Frame::ConnectAck {
flags: 0,
selected_version: ProtocolVersion::new(1, 0),
capabilities: 0,
},
Frame::SubscribeAck {
flags: 0,
stream_id: 1,
subscription_id: 7,
selected_schema: SchemaId::new([0_u8; SchemaId::WIRE_LEN]),
},
];
for reply in &replies {
match socket.read().map_err(|error| server_error(&error))? {
Message::Binary(_) => {}
other => return Err(server_error(&format!("expected binary, read {other:?}"))),
}
socket
.send(Message::Binary(encoded(reply)?.into()))
.map_err(|error| server_error(&error))?;
}
if let Some(gate) = gate {
gate.recv_timeout(GUARD)
.map_err(|error| server_error(&error))?;
}
for bytes in after {
socket
.send(Message::Binary(bytes.into()))
.map_err(|error| server_error(&error))?;
}
loop {
match socket.read() {
Ok(Message::Close(_))
| Err(
tungstenite::Error::ConnectionClosed
| tungstenite::Error::AlreadyClosed
| tungstenite::Error::Protocol(
tungstenite::error::ProtocolError::ResetWithoutClosingHandshake,
),
) => {
return Ok(());
}
Ok(_) => {}
Err(error) => return Err(server_error(&error)),
}
}
})
}
fn open_against(
gate: Option<mpsc::Receiver<()>>,
after: Vec<Vec<u8>>,
) -> Result<
(
WebSocketSubscriptionStream,
std::thread::JoinHandle<Result<(), SdkError>>,
),
SdkError,
> {
open_limited(MailboxLimits::SUBSCRIPTION, gate, after)
}
fn open_limited(
limits: MailboxLimits,
gate: Option<mpsc::Receiver<()>>,
after: Vec<Vec<u8>>,
) -> Result<
(
WebSocketSubscriptionStream,
std::thread::JoinHandle<Result<(), SdkError>>,
),
SdkError,
> {
let listener = TcpListener::bind("127.0.0.1:0").map_err(|error| server_error(&error))?;
let address = listener
.local_addr()
.map_err(|error| server_error(&error))?;
let server = serve(listener, gate, after);
let stream = WebSocketSubscriptionStream::open_with_limits(
&format!("ws://{address}/liminal"),
"events",
vec![SchemaId::new([0_u8; SchemaId::WIRE_LEN])],
&[],
limits,
)?;
Ok((stream, server))
}
fn join(server: std::thread::JoinHandle<Result<(), SdkError>>) -> Result<(), SdkError> {
server
.join()
.map_err(|_| server_error(&"server thread panicked"))?
}
fn first_receive_after(bad: Vec<u8>) -> Result<Result<(), SdkError>, SdkError> {
let (stream, server) = open_against(None, vec![bad])?;
let received = stream.recv_timeout(GUARD).map(|_| ());
drop(stream);
join(server)?;
Ok(received)
}
fn assert_decode_error(received: &Result<(), SdkError>) {
assert!(
matches!(
received,
Err(SdkError::Protocol { description })
if description.starts_with("subscription wire codec error")
|| description.starts_with("subscription decode consumed")
),
"expected the decode error itself, received {received:?}"
);
}
#[test]
fn a_malformed_delivery_reaches_the_consumer_as_its_decode_error() -> Result<(), SdkError> {
assert_decode_error(&first_receive_after(malformed(0x19))?);
Ok(())
}
#[test]
fn a_malformed_control_frame_reaches_the_consumer_as_its_decode_error() -> Result<(), SdkError> {
assert_decode_error(&first_receive_after(malformed(0x02))?);
Ok(())
}
#[test]
fn a_stray_frame_is_counted_and_the_frame_behind_it_still_arrives() -> Result<(), SdkError> {
let pong = encoded(&Frame::Pong { flags: 0 })?;
let (stream, server) = open_against(None, vec![pong, malformed(0x19)])?;
assert_decode_error(&stream.recv_timeout(GUARD).map(|_| ()));
assert_eq!(stream.stats()?.ignored_frames, 1);
drop(stream);
join(server)
}
#[test]
fn a_subscribe_error_the_consumer_was_gone_for_is_counted() -> Result<(), SdkError> {
let (go, gate) = mpsc::channel();
let subscribe_error = encoded(&Frame::SubscribeError {
flags: 0,
stream_id: 1,
reason_code: 2,
message: Some("shed".to_string()),
})?;
let (mut stream, server) = open_against(Some(gate), vec![subscribe_error])?;
stream.mailbox.release();
go.send(()).map_err(|error| server_error(&error))?;
let reader = stream
.reader
.take()
.ok_or_else(|| server_error(&"the stream has no reader"))?;
reader
.join()
.map_err(|_| server_error(&"reader thread panicked"))?;
assert_eq!(stream.stats()?.undelivered_errors, 1);
drop(stream);
join(server)
}
fn deliver(seq: u64, payload: usize) -> Result<Vec<u8>, SdkError> {
let frame = Frame::new_deliver(
1,
seq,
liminal::protocol::MessageEnvelope::new(
SchemaId::new([0_u8; SchemaId::WIRE_LEN]),
liminal::protocol::CausalContext::independent(),
vec![0xA5; payload],
),
)
.map_err(|error| SdkError::Protocol {
description: error.to_string(),
})?;
encoded(&frame)
}
#[test]
fn a_stalled_consumer_holds_at_most_the_frame_bound_and_nothing_is_dropped() -> Result<(), SdkError>
{
let limits = MailboxLimits {
frames: 4,
bytes: 1024 * 1024,
};
let deliveries = (1..=8)
.map(|seq| deliver(seq, 1024))
.collect::<Result<Vec<_>, _>>()?;
let (stream, server) = open_limited(limits, None, deliveries)?;
let stalled = stream
.mailbox
.wait_until(GUARD, |stats, _| stats.backpressure_waits >= 1)?;
assert!(stalled, "the reader never stopped on a full mailbox");
let held = stream.stats()?;
assert_eq!(held.held_frames, 4);
assert_eq!(held.held_bytes, 4 * 1024);
for expected in 1..=8 {
assert_eq!(stream.recv_timeout(GUARD)?.delivery_seq(), expected);
}
assert_eq!(stream.stats()?.peak_frames, 4);
drop(stream);
join(server)
}