use std::{
collections::{HashMap, HashSet},
sync::{
Arc, Mutex, Weak,
atomic::{AtomicBool, Ordering},
},
};
use futures::{StreamExt, future::BoxFuture, stream::FuturesUnordered};
use snafu::ResultExt;
use tokio::sync::{mpsc, watch};
use tracing::Instrument;
use super::{
WebTransportSessionId, WebTransportStreamCount,
error::{
CloseReason, RegisterSessionError, SessionCloseReason, SessionClosed, SessionDrain,
SessionFlowControlError, session_flow_control_error,
},
session::{
RoutedBiStream, RoutedUniStream,
stream::{TrackedStreamReader, TrackedStreamWriter},
},
};
use crate::{
error::Code,
quic::{ResetStreamExt, StopStreamExt},
stream_id::StreamId,
varint::VarInt,
};
const SESSION_STREAM_CHANNEL_SIZE: usize = 16;
#[derive(Debug, Default)]
struct RegistryInner {
active: HashMap<WebTransportSessionId, Weak<SessionState>>,
closed: HashSet<WebTransportSessionId>,
}
#[derive(Debug, Default, Clone)]
pub(super) struct Registry {
inner: Arc<Mutex<RegistryInner>>,
}
pub(super) enum RouteBiError {
Unknown(RoutedBiStream),
Closed(RoutedBiStream),
FlowControl(RoutedBiStream),
Rejected(RoutedBiStream),
}
pub(super) enum RouteUniError {
Unknown(RoutedUniStream),
Closed(RoutedUniStream),
FlowControl(RoutedUniStream),
Rejected(RoutedUniStream),
}
impl Registry {
pub(super) fn register(
&self,
session_id: WebTransportSessionId,
) -> Result<RegisteredSession, RegisterSessionError> {
let default_credit = default_initial_stream_credit();
self.register_with_credit(session_id, default_credit, default_credit)
}
pub(super) fn register_with_credit(
&self,
session_id: WebTransportSessionId,
bidi_credit: WebTransportStreamCount,
uni_credit: WebTransportStreamCount,
) -> Result<RegisteredSession, RegisterSessionError> {
let bidi_queue_capacity = incoming_stream_queue_capacity(bidi_credit);
let uni_queue_capacity = incoming_stream_queue_capacity(uni_credit);
let (bidi_tx, bidi_rx) = mpsc::channel(bidi_queue_capacity);
let (uni_tx, uni_rx) = mpsc::channel(uni_queue_capacity);
let Ok(mut inner) = self.inner.lock() else {
return Err(RegisterSessionError::RegistryPoisoned);
};
if inner.active.contains_key(&session_id) || inner.closed.contains(&session_id) {
return Err(RegisterSessionError::AlreadyRegistered { session_id });
}
let state = Arc::new(SessionState::new(
session_id,
self.clone(),
bidi_tx,
uni_tx,
bidi_credit,
uni_credit,
bidi_queue_capacity,
uni_queue_capacity,
));
inner.active.insert(session_id, Arc::downgrade(&state));
Ok(RegisteredSession {
state,
bidi_rx,
uni_rx,
})
}
pub(super) fn unregister(&self, session_id: WebTransportSessionId) {
self.close(session_id);
}
fn close(&self, session_id: WebTransportSessionId) {
let Ok(mut inner) = self.inner.lock() else {
tracing::debug!(?session_id, "webtransport session registry lock poisoned");
return;
};
inner.active.remove(&session_id);
inner.closed.insert(session_id);
}
pub(super) fn route_bi(
&self,
session_id: WebTransportSessionId,
stream: RoutedBiStream,
) -> Result<(), RouteBiError> {
let state = {
let Ok(inner) = self.inner.lock() else {
tracing::debug!(session_id = %session_id, "webtransport session registry lock poisoned");
return Err(RouteBiError::Rejected(stream));
};
let Some(state) = inner
.active
.get(&session_id)
.and_then(std::sync::Weak::upgrade)
else {
if inner.closed.contains(&session_id) {
tracing::debug!(session_id = %session_id, "webtransport bidi stream belongs to closed session");
return Err(RouteBiError::Closed(stream));
}
tracing::debug!(session_id = %session_id, "no registered session for webtransport bidi stream");
return Err(RouteBiError::Unknown(stream));
};
state
};
state.route_incoming_bi(stream)
}
pub(super) fn route_uni(
&self,
session_id: WebTransportSessionId,
stream: RoutedUniStream,
) -> Result<(), RouteUniError> {
let state = {
let Ok(inner) = self.inner.lock() else {
tracing::debug!(session_id = %session_id, "webtransport session registry lock poisoned");
return Err(RouteUniError::Rejected(stream));
};
let Some(state) = inner
.active
.get(&session_id)
.and_then(std::sync::Weak::upgrade)
else {
if inner.closed.contains(&session_id) {
tracing::debug!(session_id = %session_id, "webtransport uni stream belongs to closed session");
return Err(RouteUniError::Closed(stream));
}
tracing::debug!(session_id = %session_id, "no registered session for webtransport uni stream");
return Err(RouteUniError::Unknown(stream));
};
state
};
state.route_incoming_uni(stream)
}
}
impl Registry {
pub(super) fn len(&self) -> usize {
self.inner
.lock()
.map(|inner| inner.active.len())
.unwrap_or(0)
}
}
#[derive(Debug, Clone, Copy)]
struct IncomingStreamCredit {
advertised_max: WebTransportStreamCount,
received: WebTransportStreamCount,
queued: usize,
}
impl IncomingStreamCredit {
const fn new(advertised_max: WebTransportStreamCount) -> Self {
Self {
advertised_max,
received: WebTransportStreamCount::ZERO,
queued: 0,
}
}
fn reserve_incoming(&mut self, queue_capacity: usize) -> Result<(), SessionFlowControlError> {
if self.received >= self.advertised_max {
return Err(SessionFlowControlError::ExceededStreamCredit);
}
if self.queued >= queue_capacity {
return Err(SessionFlowControlError::QueueCapacityInvariant);
}
self.received = self
.received
.checked_increment()
.context(session_flow_control_error::StreamCountSnafu)?;
self.queued += 1;
Ok(())
}
fn accept_one(&mut self) -> Result<WebTransportStreamCount, SessionFlowControlError> {
self.queued = self.queued.saturating_sub(1);
self.advertised_max = self
.advertised_max
.checked_increment()
.context(session_flow_control_error::StreamCountSnafu)?;
Ok(self.advertised_max)
}
}
#[derive(Debug)]
struct LocalOpenCredit {
peer_max: WebTransportStreamCount,
opened: WebTransportStreamCount,
last_blocked_sent: Option<WebTransportStreamCount>,
changed: watch::Sender<WebTransportStreamCount>,
}
impl LocalOpenCredit {
fn new(peer_max: WebTransportStreamCount) -> Self {
let (changed, _rx) = watch::channel(peer_max);
Self {
peer_max,
opened: WebTransportStreamCount::ZERO,
last_blocked_sent: None,
changed,
}
}
fn try_reserve(&mut self) -> Result<(), WebTransportStreamCount> {
if self.opened >= self.peer_max {
return Err(self.peer_max);
}
self.opened = self
.opened
.checked_increment()
.expect("opened stream count cannot overflow below peer maximum");
Ok(())
}
fn block(&mut self) -> LocalStreamCreditBlock {
let maximum = self.peer_max;
let send_blocked = self.last_blocked_sent != Some(maximum);
if send_blocked {
self.last_blocked_sent = Some(maximum);
}
LocalStreamCreditBlock {
maximum,
send_blocked,
changed: self.changed.subscribe(),
}
}
fn update_peer_max(
&mut self,
peer_max: WebTransportStreamCount,
) -> Result<(), SessionFlowControlError> {
if peer_max < self.peer_max {
return Err(SessionFlowControlError::DecreasingMaxStreams);
}
if peer_max > self.peer_max {
self.peer_max = peer_max;
self.last_blocked_sent = None;
let _ = self.changed.send(peer_max);
}
Ok(())
}
}
pub(super) struct LocalStreamCreditBlock {
pub(super) maximum: WebTransportStreamCount,
pub(super) send_blocked: bool,
pub(super) changed: watch::Receiver<WebTransportStreamCount>,
}
pub(super) enum LocalStreamCreditReservation {
Reserved,
Blocked(LocalStreamCreditBlock),
}
#[derive(Debug)]
pub(super) struct RegisteredSession {
pub(super) state: Arc<SessionState>,
pub(super) bidi_rx: mpsc::Receiver<RoutedBiStream>,
pub(super) uni_rx: mpsc::Receiver<RoutedUniStream>,
}
#[derive(Debug)]
pub(super) struct SessionState {
session_id: WebTransportSessionId,
registry: Registry,
closed: AtomicBool,
close_reason: watch::Sender<Option<CloseReason>>,
drain_status: watch::Sender<Option<SessionDrain>>,
bidi_tx: mpsc::Sender<RoutedBiStream>,
uni_tx: mpsc::Sender<RoutedUniStream>,
bidi_credit: Mutex<IncomingStreamCredit>,
uni_credit: Mutex<IncomingStreamCredit>,
local_bidi_credit: Mutex<LocalOpenCredit>,
local_uni_credit: Mutex<LocalOpenCredit>,
bidi_queue_capacity: usize,
uni_queue_capacity: usize,
tracked_readers: Mutex<HashMap<StreamId, TrackedStreamReader>>,
tracked_writers: Mutex<HashMap<StreamId, TrackedStreamWriter>>,
}
impl SessionState {
#[allow(clippy::too_many_arguments)]
fn new(
session_id: WebTransportSessionId,
registry: Registry,
bidi_tx: mpsc::Sender<RoutedBiStream>,
uni_tx: mpsc::Sender<RoutedUniStream>,
bidi_credit: WebTransportStreamCount,
uni_credit: WebTransportStreamCount,
bidi_queue_capacity: usize,
uni_queue_capacity: usize,
) -> Self {
let (close_reason, _close_rx) = watch::channel(None);
let (drain_status, _drain_rx) = watch::channel(None);
let default_local_credit = default_initial_stream_credit();
Self {
session_id,
registry,
closed: AtomicBool::new(false),
close_reason,
drain_status,
bidi_tx,
uni_tx,
bidi_credit: Mutex::new(IncomingStreamCredit::new(bidi_credit)),
uni_credit: Mutex::new(IncomingStreamCredit::new(uni_credit)),
local_bidi_credit: Mutex::new(LocalOpenCredit::new(default_local_credit)),
local_uni_credit: Mutex::new(LocalOpenCredit::new(default_local_credit)),
bidi_queue_capacity,
uni_queue_capacity,
tracked_readers: Mutex::new(HashMap::new()),
tracked_writers: Mutex::new(HashMap::new()),
}
}
pub(super) fn id(&self) -> WebTransportSessionId {
self.session_id
}
pub(super) fn check_open(&self) -> Result<(), SessionClosed> {
if self.closed.load(Ordering::Acquire) {
Err(SessionClosed)
} else {
Ok(())
}
}
pub(super) fn set_local_stream_credit(
&self,
peer_bidi_credit: WebTransportStreamCount,
peer_uni_credit: WebTransportStreamCount,
) {
if let Ok(mut credit) = self.local_bidi_credit.lock() {
*credit = LocalOpenCredit::new(peer_bidi_credit);
} else {
tracing::debug!(
session_id = %self.session_id,
"webtransport session local bidi credit lock poisoned"
);
self.close();
}
if let Ok(mut credit) = self.local_uni_credit.lock() {
*credit = LocalOpenCredit::new(peer_uni_credit);
} else {
tracing::debug!(
session_id = %self.session_id,
"webtransport session local uni credit lock poisoned"
);
self.close();
}
}
pub(super) fn close(&self) {
self.close_with_reason(CloseReason::Session(SessionCloseReason::ControlStreamError));
}
pub(super) fn close_with_reason(&self, reason: CloseReason) {
if !self.closed.swap(true, Ordering::AcqRel) {
let _ = self.close_reason.send(Some(reason.clone()));
let _ = self.drain_status.send(Some(SessionDrain::Closed(reason)));
self.registry.unregister(self.session_id);
let (readers, writers) = self.take_tracked_streams();
spawn_tracked_stream_cleanup(self.session_id, readers, writers);
}
}
pub(super) fn drain_with_reason(&self, drain: SessionDrain) {
if self.drain_status.borrow().is_none() {
let _ = self.drain_status.send(Some(drain));
}
}
pub(super) async fn closed(&self) -> CloseReason {
let mut reason = self.close_reason.subscribe();
loop {
if let Some(reason) = reason.borrow().clone() {
return reason;
}
if reason.changed().await.is_err() {
return CloseReason::Session(SessionCloseReason::ControlStreamError);
}
}
}
pub(super) async fn drained(&self) -> SessionDrain {
let mut drain = self.drain_status.subscribe();
loop {
if let Some(drain) = drain.borrow().clone() {
return drain;
}
if drain.changed().await.is_err() {
return SessionDrain::Closed(CloseReason::Session(
SessionCloseReason::ControlStreamError,
));
}
}
}
pub(super) fn insert_tracked_bi(
&self,
stream_id: StreamId,
reader: TrackedStreamReader,
writer: TrackedStreamWriter,
) -> Result<(), SessionClosed> {
self.check_open()?;
let Ok(mut readers) = self.tracked_readers.lock() else {
tracing::debug!(
session_id = %self.session_id,
"webtransport session reader tracking lock poisoned"
);
return Err(SessionClosed);
};
let Ok(mut writers) = self.tracked_writers.lock() else {
tracing::debug!(
session_id = %self.session_id,
"webtransport session writer tracking lock poisoned"
);
return Err(SessionClosed);
};
if self.closed.load(Ordering::Acquire) {
return Err(SessionClosed);
}
readers.insert(stream_id, reader);
writers.insert(stream_id, writer);
if self.closed.load(Ordering::Acquire) {
readers.remove(&stream_id);
writers.remove(&stream_id);
Err(SessionClosed)
} else {
Ok(())
}
}
pub(super) fn insert_tracked_reader(
&self,
stream_id: StreamId,
reader: TrackedStreamReader,
) -> Result<(), SessionClosed> {
self.check_open()?;
let Ok(mut readers) = self.tracked_readers.lock() else {
tracing::debug!(
session_id = %self.session_id,
"webtransport session reader tracking lock poisoned"
);
return Err(SessionClosed);
};
if self.closed.load(Ordering::Acquire) {
return Err(SessionClosed);
}
readers.insert(stream_id, reader);
if self.closed.load(Ordering::Acquire) {
readers.remove(&stream_id);
Err(SessionClosed)
} else {
Ok(())
}
}
pub(super) fn insert_tracked_writer(
&self,
stream_id: StreamId,
writer: TrackedStreamWriter,
) -> Result<(), SessionClosed> {
self.check_open()?;
let Ok(mut writers) = self.tracked_writers.lock() else {
tracing::debug!(
session_id = %self.session_id,
"webtransport session writer tracking lock poisoned"
);
return Err(SessionClosed);
};
if self.closed.load(Ordering::Acquire) {
return Err(SessionClosed);
}
writers.insert(stream_id, writer);
if self.closed.load(Ordering::Acquire) {
writers.remove(&stream_id);
Err(SessionClosed)
} else {
Ok(())
}
}
pub(super) fn route_incoming_bi(&self, stream: RoutedBiStream) -> Result<(), RouteBiError> {
if self.check_open().is_err() {
return Err(RouteBiError::Closed(stream));
}
let Ok(mut credit) = self.bidi_credit.lock() else {
tracing::debug!(
session_id = %self.session_id,
"webtransport session bidi credit lock poisoned"
);
self.close();
return Err(RouteBiError::Rejected(stream));
};
if let Err(error) = credit.reserve_incoming(self.bidi_queue_capacity) {
let report = snafu::Report::from_error(&error);
tracing::debug!(
session_id = %self.session_id,
error = %report,
"webtransport session bidi stream credit exhausted"
);
drop(credit);
self.close();
return Err(RouteBiError::FlowControl(stream));
}
drop(credit);
match self.bidi_tx.try_send(stream) {
Ok(()) => Ok(()),
Err(error) => {
tracing::debug!(
session_id = %self.session_id,
"session bidi channel full or closed, rejecting stream"
);
self.close();
Err(RouteBiError::Rejected(error.into_inner()))
}
}
}
pub(super) fn route_incoming_uni(&self, stream: RoutedUniStream) -> Result<(), RouteUniError> {
if self.check_open().is_err() {
return Err(RouteUniError::Closed(stream));
}
let Ok(mut credit) = self.uni_credit.lock() else {
tracing::debug!(
session_id = %self.session_id,
"webtransport session uni credit lock poisoned"
);
self.close();
return Err(RouteUniError::Rejected(stream));
};
if let Err(error) = credit.reserve_incoming(self.uni_queue_capacity) {
let report = snafu::Report::from_error(&error);
tracing::debug!(
session_id = %self.session_id,
error = %report,
"webtransport session uni stream credit exhausted"
);
drop(credit);
self.close();
return Err(RouteUniError::FlowControl(stream));
}
drop(credit);
match self.uni_tx.try_send(stream) {
Ok(()) => Ok(()),
Err(error) => {
tracing::debug!(
session_id = %self.session_id,
"session uni channel full or closed, rejecting stream"
);
self.close();
Err(RouteUniError::Rejected(error.into_inner()))
}
}
}
pub(super) fn accept_incoming_bi(
&self,
) -> Result<WebTransportStreamCount, SessionFlowControlError> {
let Ok(mut credit) = self.bidi_credit.lock() else {
tracing::debug!(
session_id = %self.session_id,
"webtransport session bidi credit lock poisoned"
);
return Err(SessionFlowControlError::QueueCapacityInvariant);
};
credit.accept_one()
}
pub(super) fn accept_incoming_uni(
&self,
) -> Result<WebTransportStreamCount, SessionFlowControlError> {
let Ok(mut credit) = self.uni_credit.lock() else {
tracing::debug!(
session_id = %self.session_id,
"webtransport session uni credit lock poisoned"
);
return Err(SessionFlowControlError::QueueCapacityInvariant);
};
credit.accept_one()
}
pub(super) fn reserve_local_bidi(&self) -> Result<LocalStreamCreditReservation, SessionClosed> {
self.reserve_local_credit(&self.local_bidi_credit, "bidi")
}
pub(super) fn reserve_local_uni(&self) -> Result<LocalStreamCreditReservation, SessionClosed> {
self.reserve_local_credit(&self.local_uni_credit, "uni")
}
fn reserve_local_credit(
&self,
credit: &Mutex<LocalOpenCredit>,
direction: &'static str,
) -> Result<LocalStreamCreditReservation, SessionClosed> {
self.check_open()?;
let Ok(mut credit) = credit.lock() else {
tracing::debug!(
session_id = %self.session_id,
direction,
"webtransport session local stream credit lock poisoned"
);
self.close();
return Err(SessionClosed);
};
match credit.try_reserve() {
Ok(()) => Ok(LocalStreamCreditReservation::Reserved),
Err(_) => Ok(LocalStreamCreditReservation::Blocked(credit.block())),
}
}
pub(super) fn update_peer_bidi_max(
&self,
peer_max: WebTransportStreamCount,
) -> Result<(), SessionFlowControlError> {
self.update_peer_max(&self.local_bidi_credit, peer_max, "bidi")
}
pub(super) fn update_peer_uni_max(
&self,
peer_max: WebTransportStreamCount,
) -> Result<(), SessionFlowControlError> {
self.update_peer_max(&self.local_uni_credit, peer_max, "uni")
}
fn update_peer_max(
&self,
credit: &Mutex<LocalOpenCredit>,
peer_max: WebTransportStreamCount,
direction: &'static str,
) -> Result<(), SessionFlowControlError> {
let Ok(mut credit) = credit.lock() else {
tracing::debug!(
session_id = %self.session_id,
direction,
"webtransport session local stream credit lock poisoned"
);
self.close();
return Err(SessionFlowControlError::QueueCapacityInvariant);
};
credit.update_peer_max(peer_max)
}
pub(super) fn remove_tracked_reader(&self, stream_id: StreamId) {
let Ok(mut readers) = self.tracked_readers.lock() else {
tracing::debug!(
session_id = %self.session_id,
"webtransport session reader tracking lock poisoned"
);
return;
};
readers.remove(&stream_id);
}
pub(super) fn remove_tracked_writer(&self, stream_id: StreamId) {
let Ok(mut writers) = self.tracked_writers.lock() else {
tracing::debug!(
session_id = %self.session_id,
"webtransport session writer tracking lock poisoned"
);
return;
};
writers.remove(&stream_id);
}
fn take_tracked_streams(&self) -> (Vec<TrackedStreamReader>, Vec<TrackedStreamWriter>) {
let readers = match self.tracked_readers.lock() {
Ok(mut readers) => readers.drain().map(|(_stream_id, reader)| reader).collect(),
Err(_) => {
tracing::debug!(
session_id = %self.session_id,
"webtransport session reader tracking lock poisoned"
);
Vec::new()
}
};
let writers = match self.tracked_writers.lock() {
Ok(mut writers) => writers.drain().map(|(_stream_id, writer)| writer).collect(),
Err(_) => {
tracing::debug!(
session_id = %self.session_id,
"webtransport session writer tracking lock poisoned"
);
Vec::new()
}
};
(readers, writers)
}
}
fn spawn_tracked_stream_cleanup(
session_id: WebTransportSessionId,
readers: Vec<TrackedStreamReader>,
writers: Vec<TrackedStreamWriter>,
) {
if readers.is_empty() && writers.is_empty() {
return;
}
let cleanup = cleanup_tracked_streams(session_id, readers, writers);
match tokio::runtime::Handle::try_current() {
Ok(handle) => {
let _cleanup_task = handle.spawn(cleanup.in_current_span());
}
Err(error) => {
let report = snafu::Report::from_error(&error);
tracing::debug!(
session_id = %session_id,
error = %report,
"failed to spawn webtransport session stream cleanup"
);
}
}
}
async fn cleanup_tracked_streams(
session_id: WebTransportSessionId,
readers: Vec<TrackedStreamReader>,
writers: Vec<TrackedStreamWriter>,
) {
let mut cleanup = FuturesUnordered::<BoxFuture<'static, ()>>::new();
for mut reader in readers {
cleanup.push(Box::pin(async move {
if let Err(error) = reader.stop(Code::WT_SESSION_GONE.into_inner()).await {
let report = snafu::Report::from_error(&error);
tracing::debug!(
session_id = %session_id,
error = %report,
"failed to stop webtransport session stream reader"
);
}
}));
}
for mut writer in writers {
cleanup.push(Box::pin(async move {
if let Err(error) = writer.reset(Code::WT_SESSION_GONE.into_inner()).await {
let report = snafu::Report::from_error(&error);
tracing::debug!(
session_id = %session_id,
error = %report,
"failed to reset webtransport session stream writer"
);
}
}));
}
while cleanup.next().await.is_some() {}
}
impl Drop for SessionState {
fn drop(&mut self) {
self.close();
}
}
fn default_initial_stream_credit() -> WebTransportStreamCount {
WebTransportStreamCount::try_from(VarInt::from_u32(SESSION_STREAM_CHANNEL_SIZE as u32))
.expect("default webtransport stream credit is valid")
}
fn incoming_stream_queue_capacity(credit: WebTransportStreamCount) -> usize {
usize::try_from(credit.into_varint().into_inner())
.unwrap_or(SESSION_STREAM_CHANNEL_SIZE)
.clamp(1, SESSION_STREAM_CHANNEL_SIZE)
}
#[cfg(test)]
mod tests {
use std::{
collections::VecDeque,
panic::{AssertUnwindSafe, catch_unwind},
pin::Pin,
sync::{Arc, Mutex},
task::{Context, Poll},
};
use bytes::Bytes;
use futures::{Sink, SinkExt, Stream, StreamExt};
use super::*;
use crate::{
quic::{
self, BoxQuicStreamReader, BoxQuicStreamWriter, GetStreamIdExt, ResetStreamExt,
StopStreamExt,
},
varint::VarInt,
webtransport::WebTransportStreamCount,
};
#[derive(Debug, Default)]
struct StreamState {
written: Mutex<Vec<u8>>,
}
#[derive(Debug)]
struct TestReadStream {
chunks: VecDeque<Bytes>,
stream_id: VarInt,
}
impl Stream for TestReadStream {
type Item = Result<Bytes, quic::StreamError>;
fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Poll::Ready(self.chunks.pop_front().map(Ok))
}
}
impl quic::GetStreamId for TestReadStream {
fn poll_stream_id(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Result<VarInt, quic::StreamError>> {
Poll::Ready(Ok(self.stream_id))
}
}
impl quic::StopStream for TestReadStream {
fn poll_stop(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
_code: VarInt,
) -> Poll<Result<(), quic::StreamError>> {
Poll::Ready(Ok(()))
}
}
#[derive(Debug)]
struct TestWriteStream {
state: Arc<StreamState>,
stream_id: VarInt,
}
impl Sink<Bytes> for TestWriteStream {
type Error = quic::StreamError;
fn poll_ready(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn start_send(self: Pin<&mut Self>, item: Bytes) -> Result<(), Self::Error> {
self.state
.written
.lock()
.expect("written lock poisoned")
.extend_from_slice(&item);
Ok(())
}
fn poll_flush(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn poll_close(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
}
impl quic::GetStreamId for TestWriteStream {
fn poll_stream_id(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Result<VarInt, quic::StreamError>> {
Poll::Ready(Ok(self.stream_id))
}
}
impl quic::ResetStream for TestWriteStream {
fn poll_reset(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
_code: VarInt,
) -> Poll<Result<(), quic::StreamError>> {
Poll::Ready(Ok(()))
}
}
fn test_read_stream(id: u32, bytes: Vec<u8>) -> BoxQuicStreamReader {
Box::pin(TestReadStream {
chunks: VecDeque::from([Bytes::from(bytes)]),
stream_id: VarInt::from_u32(id),
}) as BoxQuicStreamReader
}
fn test_write_stream(id: u32, state: Arc<StreamState>) -> BoxQuicStreamWriter {
Box::pin(TestWriteStream {
state,
stream_id: VarInt::from_u32(id),
}) as BoxQuicStreamWriter
}
fn bidi_stream(id: u32) -> RoutedBiStream {
let state = Arc::new(StreamState::default());
(
test_read_stream(id, vec![id as u8]),
test_write_stream(id, Arc::clone(&state)),
)
}
fn uni_stream(id: u32) -> RoutedUniStream {
test_read_stream(id, vec![id as u8])
}
fn wt_session_id(id: u32) -> WebTransportSessionId {
WebTransportSessionId::try_from(StreamId::from(VarInt::from_u32(id)))
.expect("test id must be a valid webtransport session id")
}
fn poison_registry(registry: &Registry) {
let registry = registry.clone();
let _ = catch_unwind(AssertUnwindSafe(move || {
let _guard = registry.inner.lock().expect("registry lock should succeed");
panic!("poison registry mutex");
}));
}
fn route_bi_error_stream(error: RouteBiError) -> RoutedBiStream {
match error {
RouteBiError::Unknown(stream)
| RouteBiError::Closed(stream)
| RouteBiError::FlowControl(stream)
| RouteBiError::Rejected(stream) => stream,
}
}
fn route_uni_error_stream(error: RouteUniError) -> RoutedUniStream {
match error {
RouteUniError::Unknown(stream)
| RouteUniError::Closed(stream)
| RouteUniError::FlowControl(stream)
| RouteUniError::Rejected(stream) => stream,
}
}
#[tokio::test]
async fn test_read_stream_reads_chunks_and_stops() {
let mut stream = test_read_stream(40, b"chunk".to_vec());
assert_eq!(
stream
.next()
.await
.expect("read stream should yield one chunk")
.expect("read chunk should succeed"),
Bytes::from_static(b"chunk")
);
assert!(
stream.next().await.is_none(),
"read stream should be exhausted"
);
stream
.stop(VarInt::from_u32(41))
.await
.expect("read stream stop should succeed");
}
#[tokio::test]
async fn test_write_stream_writes_flushes_closes_and_resets() {
let state = Arc::new(StreamState::default());
let mut stream = test_write_stream(42, Arc::clone(&state));
assert_eq!(
stream
.stream_id()
.await
.expect("write stream id should be readable"),
VarInt::from_u32(42)
);
stream
.send(Bytes::from_static(b"payload"))
.await
.expect("write stream send should succeed");
stream
.close()
.await
.expect("write stream close should succeed");
stream
.reset(VarInt::from_u32(43))
.await
.expect("write stream reset should succeed");
assert_eq!(
*state
.written
.lock()
.expect("written lock should not poison"),
b"payload".to_vec()
);
}
#[test]
fn register_len_duplicate_and_close_unregister() {
let registry = Registry::default();
let session_id = wt_session_id(4);
let registered = registry
.register(session_id)
.expect("first registration should succeed");
assert_eq!(registry.len(), 1);
assert_eq!(registered.state.id(), session_id);
assert!(registered.state.check_open().is_ok());
let error = registry
.register(session_id)
.expect_err("duplicate session should be rejected");
assert!(matches!(
error,
RegisterSessionError::AlreadyRegistered { session_id: duplicate } if duplicate == session_id
));
registered.state.close();
assert_eq!(registry.len(), 0);
assert!(registered.state.check_open().is_err());
registered.state.close();
assert_eq!(registry.len(), 0);
}
#[test]
fn dropping_registered_session_unregisters_it() {
let registry = Registry::default();
let session_id = wt_session_id(8);
let registered = registry
.register(session_id)
.expect("registration should succeed");
assert_eq!(registry.len(), 1);
drop(registered);
assert_eq!(registry.len(), 0);
}
#[tokio::test]
async fn route_unknown_sessions_return_original_streams() {
let registry = Registry::default();
let session_id = wt_session_id(4);
let bidi = bidi_stream(1);
let error = registry
.route_bi(session_id, bidi)
.expect_err("unknown bidi session should reject stream");
assert!(matches!(&error, RouteBiError::Unknown(_)));
let mut returned_bidi = route_bi_error_stream(error);
assert_eq!(
returned_bidi
.0
.stream_id()
.await
.expect("returned bidi stream id should be readable"),
VarInt::from_u32(1)
);
let uni = uni_stream(2);
let error = registry
.route_uni(session_id, uni)
.expect_err("unknown uni session should reject stream");
assert!(matches!(&error, RouteUniError::Unknown(_)));
let mut returned_uni = route_uni_error_stream(error);
assert_eq!(
returned_uni
.stream_id()
.await
.expect("returned uni stream id should be readable"),
VarInt::from_u32(2)
);
}
#[tokio::test]
async fn route_known_sessions_deliver_streams_to_receivers() {
let registry = Registry::default();
let session_id = wt_session_id(4);
let mut registered = registry
.register(session_id)
.expect("registration should succeed");
assert!(registry.route_bi(session_id, bidi_stream(3)).is_ok());
assert!(registry.route_uni(session_id, uni_stream(4)).is_ok());
let (mut bidi_reader, _bidi_writer) = registered
.bidi_rx
.recv()
.await
.expect("bidi receiver should get a stream");
let mut uni_reader = registered
.uni_rx
.recv()
.await
.expect("uni receiver should get a stream");
assert_eq!(
bidi_reader
.stream_id()
.await
.expect("bidi stream id should be readable"),
VarInt::from_u32(3)
);
assert_eq!(
uni_reader
.stream_id()
.await
.expect("uni stream id should be readable"),
VarInt::from_u32(4)
);
}
#[tokio::test]
async fn route_incoming_bidi_closes_session_when_peer_exceeds_advertised_credit() {
let registry = Registry::default();
let session_id = wt_session_id(4);
let mut registered = registry
.register_with_credit(
session_id,
WebTransportStreamCount::try_from(VarInt::from_u32(1)).expect("bidi credit"),
WebTransportStreamCount::try_from(VarInt::from_u32(0)).expect("uni credit"),
)
.expect("registration should succeed");
assert!(
registry.route_bi(session_id, bidi_stream(8)).is_ok(),
"first stream should be allowed"
);
let error = registry
.route_bi(session_id, bidi_stream(12))
.expect_err("second stream exceeds credit");
assert!(matches!(error, RouteBiError::FlowControl(_)));
assert!(registered.state.check_open().is_err());
let (_reader, _writer) = registered
.bidi_rx
.recv()
.await
.expect("first stream remains queued");
}
#[tokio::test]
async fn route_incoming_uni_closes_session_when_peer_exceeds_advertised_credit() {
let registry = Registry::default();
let session_id = wt_session_id(4);
let mut registered = registry
.register_with_credit(
session_id,
WebTransportStreamCount::try_from(VarInt::from_u32(0)).expect("bidi credit"),
WebTransportStreamCount::try_from(VarInt::from_u32(1)).expect("uni credit"),
)
.expect("registration should succeed");
assert!(
registry.route_uni(session_id, uni_stream(10)).is_ok(),
"first stream should be allowed"
);
let error = registry
.route_uni(session_id, uni_stream(14))
.expect_err("second stream exceeds credit");
assert!(matches!(error, RouteUniError::FlowControl(_)));
assert!(registered.state.check_open().is_err());
let _reader = registered
.uni_rx
.recv()
.await
.expect("first stream remains queued");
}
#[test]
fn route_rejects_closed_or_full_channels() {
let registry = Registry::default();
let bidi_session_id = wt_session_id(4);
let bidi_registered = registry
.register(bidi_session_id)
.expect("registration should succeed");
drop(bidi_registered.bidi_rx);
assert!(matches!(
registry.route_bi(bidi_session_id, bidi_stream(5)),
Err(RouteBiError::Rejected(_))
));
let uni_session_id = wt_session_id(8);
let uni_registered = registry
.register(uni_session_id)
.expect("registration should succeed");
drop(uni_registered.uni_rx);
assert!(matches!(
registry.route_uni(uni_session_id, uni_stream(6)),
Err(RouteUniError::Rejected(_))
));
}
#[tokio::test]
async fn route_closed_channels_return_original_streams() {
let registry = Registry::default();
let bidi_session_id = wt_session_id(16);
let bidi_registered = registry
.register(bidi_session_id)
.expect("registration should succeed");
drop(bidi_registered.bidi_rx);
let error = registry
.route_bi(bidi_session_id, bidi_stream(17))
.expect_err("closed bidi receiver should reject stream");
assert!(matches!(&error, RouteBiError::Rejected(_)));
let mut returned_bidi = route_bi_error_stream(error);
assert_eq!(
returned_bidi
.0
.stream_id()
.await
.expect("returned bidi stream id should be readable"),
VarInt::from_u32(17)
);
let uni_session_id = wt_session_id(20);
let uni_registered = registry
.register(uni_session_id)
.expect("registration should succeed");
drop(uni_registered.uni_rx);
let error = registry
.route_uni(uni_session_id, uni_stream(18))
.expect_err("closed uni receiver should reject stream");
assert!(matches!(&error, RouteUniError::Rejected(_)));
let mut returned_uni = route_uni_error_stream(error);
assert_eq!(
returned_uni
.stream_id()
.await
.expect("returned uni stream id should be readable"),
VarInt::from_u32(18)
);
}
#[test]
fn route_flow_control_when_initial_credit_is_exhausted() {
let registry = Registry::default();
let bidi_session_id = wt_session_id(24);
let _bidi_registered = registry
.register(bidi_session_id)
.expect("registration should succeed");
for id in 0..SESSION_STREAM_CHANNEL_SIZE {
assert!(
registry
.route_bi(bidi_session_id, bidi_stream(id as u32))
.is_ok()
);
}
assert!(matches!(
registry.route_bi(bidi_session_id, bidi_stream(99)),
Err(RouteBiError::FlowControl(_))
));
let uni_session_id = wt_session_id(28);
let _uni_registered = registry
.register(uni_session_id)
.expect("registration should succeed");
for id in 0..SESSION_STREAM_CHANNEL_SIZE {
assert!(
registry
.route_uni(uni_session_id, uni_stream(id as u32))
.is_ok()
);
}
assert!(matches!(
registry.route_uni(uni_session_id, uni_stream(100)),
Err(RouteUniError::FlowControl(_))
));
}
#[tokio::test]
async fn route_uses_exact_registered_session_id() {
let registry = Registry::default();
let first_session_id = wt_session_id(32);
let second_session_id = wt_session_id(36);
let mut first = registry
.register(first_session_id)
.expect("first registration should succeed");
let mut second = registry
.register(second_session_id)
.expect("second registration should succeed");
assert!(
registry
.route_bi(second_session_id, bidi_stream(37))
.is_ok()
);
assert!(
registry
.route_uni(second_session_id, uni_stream(38))
.is_ok()
);
assert!(matches!(
first.bidi_rx.try_recv(),
Err(mpsc::error::TryRecvError::Empty)
));
assert!(matches!(
first.uni_rx.try_recv(),
Err(mpsc::error::TryRecvError::Empty)
));
let (mut second_bidi_reader, _second_bidi_writer) = second
.bidi_rx
.recv()
.await
.expect("second session should receive routed bidi stream");
assert_eq!(
second_bidi_reader
.stream_id()
.await
.expect("second bidi stream id should be readable"),
VarInt::from_u32(37)
);
let mut second_uni_reader = second
.uni_rx
.recv()
.await
.expect("second session should receive routed uni stream");
assert_eq!(
second_uni_reader
.stream_id()
.await
.expect("second uni stream id should be readable"),
VarInt::from_u32(38)
);
}
#[tokio::test]
async fn route_full_channels_return_original_streams() {
let registry = Registry::default();
let bidi_session_id = wt_session_id(40);
let _bidi_registered = registry
.register(bidi_session_id)
.expect("registration should succeed");
for id in 0..SESSION_STREAM_CHANNEL_SIZE {
assert!(
registry
.route_bi(bidi_session_id, bidi_stream(id as u32))
.is_ok()
);
}
let error = registry
.route_bi(bidi_session_id, bidi_stream(21))
.expect_err("exhausted bidi credit should reject stream");
assert!(matches!(&error, RouteBiError::FlowControl(_)));
let mut returned_bidi = route_bi_error_stream(error);
assert_eq!(
returned_bidi
.0
.stream_id()
.await
.expect("returned bidi stream id should be readable"),
VarInt::from_u32(21)
);
let uni_session_id = wt_session_id(44);
let _uni_registered = registry
.register(uni_session_id)
.expect("registration should succeed");
for id in 0..SESSION_STREAM_CHANNEL_SIZE {
assert!(
registry
.route_uni(uni_session_id, uni_stream(id as u32))
.is_ok()
);
}
let error = registry
.route_uni(uni_session_id, uni_stream(22))
.expect_err("exhausted uni credit should reject stream");
assert!(matches!(&error, RouteUniError::FlowControl(_)));
let mut returned_uni = route_uni_error_stream(error);
assert_eq!(
returned_uni
.stream_id()
.await
.expect("returned uni stream id should be readable"),
VarInt::from_u32(22)
);
}
#[tokio::test]
async fn closed_session_tombstone_is_preserved_for_connection_lifetime() {
let registry = Registry::default();
let session_id = wt_session_id(24);
let registered = registry
.register(session_id)
.expect("initial registration should succeed");
registered.state.close();
assert_eq!(registry.len(), 0);
let error = registry
.route_uni(session_id, uni_stream(25))
.expect_err("closed session should no longer route streams");
assert!(matches!(&error, RouteUniError::Closed(_)));
let mut returned = route_uni_error_stream(error);
assert_eq!(
returned
.stream_id()
.await
.expect("returned uni stream id should be readable"),
VarInt::from_u32(25)
);
let error = registry
.register(session_id)
.expect_err("closed session id must not be re-registered");
assert!(matches!(
error,
RegisterSessionError::AlreadyRegistered {
session_id: duplicate
} if duplicate == session_id
));
}
#[test]
fn cloned_session_state_keeps_registration_until_last_state_drop() {
let registry = Registry::default();
let session_id = wt_session_id(28);
let registered = registry
.register(session_id)
.expect("registration should succeed");
let state = Arc::clone(®istered.state);
drop(registered);
assert_eq!(registry.len(), 1);
assert!(matches!(
registry.route_uni(session_id, uni_stream(29)),
Err(RouteUniError::Rejected(_))
));
assert!(state.check_open().is_err());
assert_eq!(registry.len(), 0);
drop(state);
assert_eq!(registry.len(), 0);
}
#[tokio::test]
async fn poisoned_registry_surfaces_errors_and_preserves_streams() {
let registry = Registry::default();
let session_id = wt_session_id(12);
poison_registry(®istry);
let error = registry
.register(session_id)
.expect_err("poisoned registry should reject new registrations");
assert!(matches!(error, RegisterSessionError::RegistryPoisoned));
assert_eq!(registry.len(), 0);
registry.unregister(session_id);
let error = registry
.route_bi(session_id, bidi_stream(13))
.expect_err("poisoned registry should return routed bidi stream");
assert!(matches!(&error, RouteBiError::Rejected(_)));
let mut returned_bidi = route_bi_error_stream(error);
assert_eq!(
returned_bidi
.0
.stream_id()
.await
.expect("returned bidi stream id should be readable"),
VarInt::from_u32(13)
);
let error = registry
.route_uni(session_id, uni_stream(14))
.expect_err("poisoned registry should return routed uni stream");
assert!(matches!(&error, RouteUniError::Rejected(_)));
let mut returned_uni = route_uni_error_stream(error);
assert_eq!(
returned_uni
.stream_id()
.await
.expect("returned uni stream id should be readable"),
VarInt::from_u32(14)
);
}
}