use std::sync::atomic::Ordering;
use std::sync::{Arc, Weak};
use std::time::Duration;
use tokio::sync::{mpsc, watch};
use weida_core::{Error, ErrorCode, Limits, LossCause, PeerIdentity};
use weida_protocol::header::{GuaranteeSet, MAX_CURSOR_RECORD_LEN};
use weida_protocol::{
Agreed, CreditHeader, CursorHeader, DataHeader, ErrorHeader, FrameKind, Hello,
MAX_PREAMBLE_LEN, Preamble, PreambleError, SubscriptionHeader, codes, decode_cursor_record,
encode_frame, negotiate, parse_preamble,
};
use crate::cursor::CursorSet;
use crate::dedup::DedupWindow;
use crate::listener::{Namespace, Route};
use crate::ordering::{GapDetector, Reassembler, Sequencer};
use crate::pubsub::SubRegistry;
use crate::runtime::{Exec, Shared};
use crate::stream::{Consumer, ConsumerId, CreditGrant, Incoming};
use crate::transfer::{IncomingMeta, IncomingRequest, IncomingTransfer};
use crate::transport::{Link, RecvHalf, SendHalf};
const CTL_QUEUE: usize = 1024;
pub(crate) enum Ctl {
SendUnsubscribe { path: Arc<str>, filter: String },
ReplyError { send: SendHalf, code: ErrorCode },
}
pub(crate) struct ConnCtx {
pub conn: Link,
pub ctl: mpsc::Sender<Ctl>,
pub limits: Limits,
pub namespace: Arc<Namespace>,
pub subs: Option<Arc<SubRegistry>>,
pub peer: Option<PeerIdentity>,
pub exec: Exec,
pub guarantees: GuaranteeSet,
pub sequencer: Sequencer,
pub gaps: GapDetector,
pub reorder: Reassembler<Held>,
pub dedup: DedupWindow,
pub parked: crate::drain::ConnDrain,
pub reports: crate::cursor::ReportTable,
pub shared: Arc<Shared>,
agreed: watch::Receiver<Option<Agreed>>,
}
pub(crate) type ConnHandle = Arc<ConnCtx>;
impl ConnCtx {
pub(crate) fn spawn(
conn: Link,
limits: Limits,
namespace: Arc<Namespace>,
subs: Option<Arc<SubRegistry>>,
exec: Exec,
guarantees: GuaranteeSet,
shared: Arc<Shared>,
) -> ConnHandle {
let (ctl_tx, ctl_rx) = mpsc::channel(CTL_QUEUE);
let (agreed_tx, agreed_rx) = watch::channel(None);
let agreed_tx = Arc::new(agreed_tx);
let streams_are_local = conn.streams_are_local();
let ctx = Arc::new(ConnCtx {
peer: conn.peer(),
conn,
ctl: ctl_tx,
limits,
namespace,
subs,
exec: exec.clone(),
guarantees,
sequencer: Sequencer::new(guarantees.ordering),
gaps: GapDetector::new(guarantees.ordering, limits.max_sequence_scopes),
reorder: Reassembler::new(
guarantees.ordering,
limits.max_reorder_hold,
limits.max_sequence_scopes,
),
dedup: DedupWindow::new(
guarantees.deduplication,
guarantees.dedup_window_ms,
limits.max_dedup_entries,
),
parked: crate::drain::ConnDrain::new(&limits, streams_are_local),
reports: crate::cursor::ReportTable::new(),
shared,
agreed: agreed_rx,
});
ctx.shared.drain.register(&ctx);
exec.spawn(driver(Arc::downgrade(&ctx), ctl_rx, exec.clone()));
exec.spawn(hello_deadline(Arc::clone(&ctx), Arc::clone(&agreed_tx)));
exec.spawn(accept_uni_loop(Arc::clone(&ctx), agreed_tx));
exec.spawn(accept_bi_loop(Arc::clone(&ctx)));
exec.spawn(send_hello(Arc::clone(&ctx), limits, guarantees));
ctx
}
pub(crate) async fn negotiated(&self) -> Result<Agreed, Error> {
let mut rx = self.agreed.clone();
loop {
if let Some(agreed) = *rx.borrow_and_update() {
return Ok(agreed);
}
tokio::select! {
changed = rx.changed() => {
if changed.is_err() {
return Err(self
.conn
.close_reason()
.unwrap_or(Error::ConnectionLost(LossCause::LocallyClosed)));
}
}
reason = self.conn.closed() => return Err(reason),
}
}
}
pub(crate) fn notify(&self, ctl: Ctl) {
if self.ctl.try_send(ctl).is_err() {
tracing::debug!("connection control queue full or closed; notification dropped");
}
}
pub(crate) async fn open_uni(&self) -> Result<SendHalf, Error> {
self.free_local_slots();
self.conn.open_uni().await
}
pub(crate) async fn open_bi(&self) -> Result<(SendHalf, RecvHalf), Error> {
self.free_local_slots();
self.conn.open_bi().await
}
fn free_local_slots(&self) {
if !self.conn.local_slots_exhausted() {
return;
}
let freed = self.parked.reap();
if freed > 0 {
tracing::debug!(
freed,
"reaped settled receipts to make room for another stream"
);
}
}
}
pub(crate) fn conn_error(e: quinn::ConnectionError) -> Error {
match e {
quinn::ConnectionError::ApplicationClosed(frame) => match frame.error_code.into_inner() {
codes::NEGOTIATION_FAILED => {
Error::Negotiation("peer closed the connection: negotiation failed".into())
}
codes::PROTOCOL_VIOLATION => {
Error::Protocol("peer closed the connection: protocol violation".into())
}
codes::LIMIT_EXCEEDED => Error::LimitExceeded,
_ => Error::ConnectionLost(LossCause::PeerClosed),
},
quinn::ConnectionError::LocallyClosed => Error::ConnectionLost(LossCause::LocallyClosed),
quinn::ConnectionError::TimedOut => Error::ConnectionLost(LossCause::IdleTimeout),
quinn::ConnectionError::Reset => Error::ConnectionLost(LossCause::Reset),
quinn::ConnectionError::TransportError(t) if is_tls_alert(t.code) => {
Error::Tls(t.to_string())
}
quinn::ConnectionError::ConnectionClosed(c) if is_tls_alert(c.error_code) => {
Error::Tls(format!("peer aborted the handshake: {c}"))
}
other => Error::Transport(other.to_string()),
}
}
fn is_tls_alert(code: quinn::TransportErrorCode) -> bool {
(0x100..0x200).contains(&u64::from(code))
}
pub(crate) fn write_error(e: quinn::WriteError) -> Error {
match e {
quinn::WriteError::Stopped(code) => codes::stop_reason(code.into_inner()).into(),
quinn::WriteError::ConnectionLost(e) => conn_error(e),
quinn::WriteError::ClosedStream => Error::Transport("stream already closed".into()),
quinn::WriteError::ZeroRttRejected => {
Error::Transport("0-RTT data rejected by the peer".into())
}
}
}
pub(crate) fn read_error(e: quinn::ReadError) -> Error {
match e {
quinn::ReadError::Reset(code) => match code.into_inner() {
codes::CANCELED => Error::Canceled,
codes::REJECTED => Error::Rejected,
codes::UNKNOWN_ENDPOINT => Error::UnknownEndpoint,
codes::UNSUPPORTED => Error::Unsupported,
other => Error::Transport(format!("peer reset the stream with code {other}")),
},
quinn::ReadError::ConnectionLost(e) => conn_error(e),
quinn::ReadError::ClosedStream => Error::Transport("stream already closed".into()),
quinn::ReadError::IllegalOrderedRead => {
Error::Transport("ordered read after unordered read".into())
}
quinn::ReadError::ZeroRttRejected => {
Error::Transport("0-RTT data rejected by the peer".into())
}
}
}
async fn driver(ctx: Weak<ConnCtx>, mut rx: mpsc::Receiver<Ctl>, exec: Exec) {
while let Some(ctl) = rx.recv().await {
let Some(ctx) = ctx.upgrade() else { break };
handle_ctl(ctx, ctl, &exec);
}
}
fn handle_ctl(ctx: ConnHandle, ctl: Ctl, exec: &Exec) {
match ctl {
Ctl::SendUnsubscribe { path, filter } => {
let header = SubscriptionHeader::new(&*path, filter).encode();
exec.spawn(async move {
if let Err(e) = write_control(&ctx.conn, FrameKind::Unsubscribe, &header).await {
tracing::debug!(error = %e, "failed to send an UNSUBSCRIBE frame");
}
});
}
Ctl::ReplyError { mut send, code } => {
exec.spawn(async move {
if let Err(e) = write_error_frame(&mut send, code).await {
tracing::debug!(error = %e, "failed to report a reply failure");
}
});
}
}
}
pub(crate) async fn write_error_frame(send: &mut SendHalf, code: ErrorCode) -> Result<(), Error> {
let frame = encode_frame(FrameKind::Error, &ErrorHeader::new(code).encode());
send.write_all(&frame).await?;
send.finish()
}
pub(crate) async fn write_control(
conn: &Link,
kind: FrameKind,
header: &[u8],
) -> Result<(), Error> {
let mut stream = if kind == FrameKind::Hello {
conn.open_control().await?
} else {
conn.open_uni().await?
};
stream.write_all(&encode_frame(kind, header)).await?;
stream.finish()?;
Ok(())
}
fn hello_for(limits: Limits, guarantees: GuaranteeSet) -> Hello {
let declaration = (!guarantees.is_core()).then_some(guarantees);
Hello {
guarantees_offered: declaration,
guarantees_required: declaration,
..Hello::v0(
limits.max_header_bytes,
u64::from(limits.max_concurrent_uni_streams),
)
}
}
async fn send_hello(ctx: ConnHandle, limits: Limits, guarantees: GuaranteeSet) {
let hello = hello_for(limits, guarantees);
if let Err(e) = write_control(&ctx.conn, FrameKind::Hello, &hello.encode()).await {
tracing::debug!(error = %e, "failed to send HELLO");
}
}
async fn hello_deadline(ctx: ConnHandle, agreed_tx: Arc<watch::Sender<Option<Agreed>>>) {
tokio::select! {
() = ctx.exec.sleep(Duration::from_millis(ctx.limits.hello_timeout_ms)) => {}
_ = ctx.conn.closed() => return,
}
if agreed_tx.borrow().is_none() {
tracing::debug!("peer HELLO did not arrive in time");
ctx.conn.close(codes::NEGOTIATION_FAILED, "hello timeout");
}
}
async fn accept_uni_loop(ctx: ConnHandle, agreed_tx: Arc<watch::Sender<Option<Agreed>>>) {
loop {
match ctx.conn.accept_uni().await {
Ok(stream) => {
if ctx.shared.drain.is_draining() {
refuse_uni(stream);
continue;
}
let ctx = Arc::clone(&ctx);
let agreed_tx = Arc::clone(&agreed_tx);
ctx.exec.clone().spawn(async move {
if let Err(e) = handle_stream(&ctx, &agreed_tx, stream).await {
tracing::debug!(error = %e, "inbound stream failed");
}
});
}
Err(e) => {
tracing::debug!(error = %e, "connection closed; uni accept loop ending");
break;
}
}
}
}
async fn accept_bi_loop(ctx: ConnHandle) {
loop {
match ctx.conn.accept_bi().await {
Ok((mut send, recv)) => {
if ctx.shared.drain.is_draining() {
refuse_uni(recv);
send.reset(codes::SHUTDOWN);
continue;
}
let ctx = Arc::clone(&ctx);
let by_path = ctx.conn.dispatch_by_path();
ctx.exec.clone().spawn(async move {
let outcome = if by_path {
handle_local(&ctx, send, recv).await
} else {
handle_bi(&ctx, send, recv).await
};
if let Err(e) = outcome {
tracing::debug!(error = %e, "inbound exchange failed");
}
});
}
Err(e) => {
tracing::debug!(error = %e, "connection closed; bidi accept loop ending");
break;
}
}
}
}
async fn handle_local(ctx: &ConnHandle, send: SendHalf, mut recv: RecvHalf) -> Result<(), Error> {
let preamble = match read_preamble(&mut recv, ctx.limits.max_header_bytes).await {
Ok(preamble) => preamble,
Err(Error::Protocol(reason)) => return violation(ctx, &reason),
Err(e) => return Err(e),
};
let header = read_header(&mut recv, &preamble).await?;
if preamble.kind != FrameKind::Hello {
ctx.negotiated().await?;
}
match preamble.kind {
FrameKind::Hello => violation(ctx, "HELLO is legal only on the control connection"),
FrameKind::Error => violation(ctx, "ERROR is legal only in answer to a request"),
FrameKind::Subscribe => handle_subscription(ctx, &header, true).await,
FrameKind::Unsubscribe => handle_subscription(ctx, &header, false).await,
FrameKind::Credit => handle_credit(ctx, &header).await,
FrameKind::Cursor => {
drop(send);
handle_cursor(ctx, recv, &header).await
}
FrameKind::Data => {
let decoded = match DataHeader::decode(&header) {
Ok(h) => h,
Err(e) => return violation(ctx, &e.to_string()),
};
let Some(path) = decoded.endpoint.clone() else {
return violation(ctx, "DATA on a local connection must name an endpoint");
};
let route = ctx.namespace.lookup(&path);
match route {
Some(Route::Request(queue)) => {
let request = IncomingRequest::new(
IncomingTransfer::new(
recv,
Arc::new(IncomingMeta::from_header(&decoded, ctx.peer.clone())),
Arc::clone(ctx),
),
send,
Arc::clone(ctx),
);
if queue.send(request).await.is_err() {
tracing::debug!(path, "endpoint went away while dispatching");
}
Ok(())
}
Some(Route::Raw(queue)) => {
let request = IncomingRequest::new(
IncomingTransfer::new(
recv,
Arc::new(IncomingMeta::from_header(&decoded, ctx.peer.clone())),
Arc::clone(ctx),
),
send,
Arc::clone(ctx),
);
if queue.send(Incoming::Exchange(request)).await.is_err() {
tracing::debug!(path, "acceptor went away while dispatching");
}
Ok(())
}
_ => {
drop(send);
handle_data(ctx, recv, &header).await
}
}
}
}
}
fn refuse_uni(mut stream: RecvHalf) {
stream.stop(codes::SHUTDOWN);
}
fn refuse_cursor(mut stream: RecvHalf) -> Result<(), Error> {
stream.stop(codes::CANCELED);
Ok(())
}
async fn handle_cursor(ctx: &ConnHandle, mut stream: RecvHalf, header: &[u8]) -> Result<(), Error> {
let head = match CursorHeader::decode(header) {
Ok(head) => head,
Err(e) => return violation(ctx, &e.to_string()),
};
let Some(tx) = ctx.reports.claim(head.report_id) else {
return refuse_cursor(stream);
};
let mut buf = [0u8; 64];
let mut pending: Vec<u8> = Vec::with_capacity(64 + MAX_CURSOR_RECORD_LEN);
let mut set = CursorSet::default();
loop {
match stream.read(&mut buf).await {
Ok(Some(0)) => continue,
Ok(Some(n)) => {
pending.extend_from_slice(&buf[..n]);
let mut at = 0;
while let Some((level, offset, used)) = match decode_cursor_record(&pending[at..]) {
Ok(record) => record,
Err(e) => return violation(ctx, &e.to_string()),
} {
at += used;
if set.advance(level, offset) {
tx.send_replace(set);
}
}
pending.drain(..at);
}
Ok(None) => {
if !pending.is_empty() {
return violation(ctx, "a cursor stream ended inside a record");
}
break;
}
Err(e) => {
tracing::debug!(error = %e, report_id = head.report_id, "cursor stream ended");
break;
}
}
}
ctx.reports.release(head.report_id);
drop(tx);
Ok(())
}
async fn read_preamble(stream: &mut RecvHalf, max_header_bytes: u64) -> Result<Preamble, Error> {
let mut scratch = [0u8; MAX_PREAMBLE_LEN];
let mut have = 0usize;
loop {
match parse_preamble(&scratch[..have], max_header_bytes) {
Ok((preamble, used)) => {
debug_assert_eq!(used, have, "the preamble is read byte by byte");
return Ok(preamble);
}
Err(PreambleError::Incomplete) => {
if have == scratch.len() {
return Err(Error::Protocol(
"preamble exceeds its maximum length".into(),
));
}
match stream.read(&mut scratch[have..have + 1]).await? {
Some(0) => continue,
Some(n) => have += n,
None => return Err(Error::Protocol("stream ended inside the preamble".into())),
}
}
Err(e) => return Err(Error::Protocol(e.to_string())),
}
}
}
async fn read_header(stream: &mut RecvHalf, preamble: &Preamble) -> Result<Vec<u8>, Error> {
let mut header = vec![0u8; preamble.header_len as usize];
stream.read_exact(&mut header).await?;
Ok(header)
}
pub(crate) async fn read_frame(
stream: &mut RecvHalf,
max_header_bytes: u64,
) -> Result<(Preamble, Vec<u8>), Error> {
let preamble = read_preamble(stream, max_header_bytes).await?;
let header = read_header(stream, &preamble).await?;
Ok((preamble, header))
}
async fn handle_stream(
ctx: &ConnHandle,
agreed_tx: &watch::Sender<Option<Agreed>>,
mut stream: RecvHalf,
) -> Result<(), Error> {
let preamble = match read_preamble(&mut stream, ctx.limits.max_header_bytes).await {
Ok(preamble) => preamble,
Err(Error::Protocol(reason)) => return violation(ctx, &reason),
Err(e) => return Err(e),
};
if preamble.kind != FrameKind::Hello {
ctx.negotiated().await?;
}
let header = read_header(&mut stream, &preamble).await?;
match preamble.kind {
FrameKind::Hello => handle_hello(ctx, agreed_tx, &header),
FrameKind::Data => handle_data(ctx, stream, &header).await,
FrameKind::Error => violation(
ctx,
"ERROR is legal only on the reply half of a bidirectional stream",
),
FrameKind::Subscribe => handle_subscription(ctx, &header, true).await,
FrameKind::Unsubscribe => handle_subscription(ctx, &header, false).await,
FrameKind::Credit => handle_credit(ctx, &header).await,
FrameKind::Cursor => handle_cursor(ctx, stream, &header).await,
}
}
fn violation(ctx: &ConnHandle, reason: &str) -> Result<(), Error> {
tracing::debug!(reason, "closing connection: protocol violation");
ctx.conn.close(codes::PROTOCOL_VIOLATION, reason);
Err(Error::Protocol(reason.to_owned()))
}
async fn handle_subscription(
ctx: &ConnHandle,
header: &[u8],
subscribe: bool,
) -> Result<(), Error> {
let header = match SubscriptionHeader::decode(header) {
Ok(h) => h,
Err(e) => return violation(ctx, &e.to_string()),
};
let Some(subs) = ctx.subs.as_ref() else {
tracing::debug!(
endpoint = %header.endpoint,
"ignoring a subscription frame: this side publishes nothing"
);
return Ok(());
};
let conn_id = ctx.conn.stable_id();
if let Some(Route::Raw(queue)) = ctx.namespace.lookup(&header.endpoint) {
if subscribe {
if subs.reserve(conn_id).is_err() {
return too_many_subscriptions(ctx);
}
ctx.namespace.note_consumer(conn_id, &header.endpoint);
let consumer = Consumer::new(
Arc::clone(ctx),
Arc::from(header.endpoint.as_str()),
header.filter,
);
if queue.send(Incoming::Subscribed(consumer)).await.is_err() {
tracing::debug!(endpoint = %header.endpoint, "acceptor went away; subscription dropped");
}
} else {
subs.release(conn_id);
ctx.namespace.forget_consumer(conn_id, &header.endpoint);
let gone = Incoming::Unsubscribed {
id: ConsumerId::from_conn(conn_id),
filter: Some(header.filter),
};
if queue.send(gone).await.is_err() {
tracing::debug!(endpoint = %header.endpoint, "acceptor went away; unsubscribe dropped");
}
}
return Ok(());
}
if subscribe {
if subs
.subscribe(&header.endpoint, ctx, header.filter)
.is_err()
{
return too_many_subscriptions(ctx);
}
} else {
subs.unsubscribe(&header.endpoint, conn_id, &header.filter);
}
Ok(())
}
fn too_many_subscriptions(ctx: &ConnHandle) -> Result<(), Error> {
tracing::debug!(
max = ctx.limits.max_subscriptions,
"subscription limit reached; closing the connection"
);
ctx.conn.close(
codes::LIMIT_EXCEEDED,
"too many subscriptions on one connection",
);
Err(Error::LimitExceeded)
}
async fn handle_credit(ctx: &ConnHandle, header: &[u8]) -> Result<(), Error> {
let header = match CreditHeader::decode(header) {
Ok(h) => h,
Err(e) => return violation(ctx, &e.to_string()),
};
let Some(Route::Raw(queue)) = ctx.namespace.lookup(&header.endpoint) else {
tracing::debug!(
endpoint = %header.endpoint,
"ignoring a credit frame: no queue serves this path"
);
return Ok(());
};
let grant = Incoming::Credit(CreditGrant {
id: ConsumerId::from_conn(ctx.conn.stable_id()),
filter: header.filter,
limit: header.limit,
});
if queue.send(grant).await.is_err() {
tracing::debug!(endpoint = %header.endpoint, "acceptor went away; credit dropped");
}
Ok(())
}
pub(crate) async fn drop_consumers(ctx: &ConnHandle) {
let conn_id = ctx.conn.stable_id();
for (path, route) in ctx.namespace.take_consumer_routes(conn_id) {
let Route::Raw(queue) = route else {
continue;
};
let gone = Incoming::Unsubscribed {
id: ConsumerId::from_conn(conn_id),
filter: None,
};
if queue.send(gone).await.is_err() {
tracing::debug!(%path, "acceptor went away before its consumer did");
}
}
}
fn handle_hello(
ctx: &ConnHandle,
agreed_tx: &watch::Sender<Option<Agreed>>,
header: &[u8],
) -> Result<(), Error> {
let theirs = match Hello::decode(header) {
Ok(h) => h,
Err(e) => return violation(ctx, &e.to_string()),
};
let ours = hello_for(ctx.limits, ctx.guarantees);
match negotiate(&ours, &theirs) {
Ok(agreed) => {
tracing::debug!(
version = agreed.version,
send_max_header_bytes = agreed.send_max_header_bytes,
"negotiated"
);
let _ = agreed_tx.send(Some(agreed));
Ok(())
}
Err(e) => {
tracing::debug!(error = %e, "negotiation failed");
ctx.conn.close(codes::NEGOTIATION_FAILED, &e.to_string());
Err(e.into())
}
}
}
async fn handle_data(ctx: &ConnHandle, stream: RecvHalf, header: &[u8]) -> Result<(), Error> {
let header = match DataHeader::decode(header) {
Ok(h) => h,
Err(e) => return violation(ctx, &e.to_string()),
};
let Some(path) = header.endpoint.clone() else {
return violation(ctx, "DATA on a unidirectional stream must name an endpoint");
};
let scope = header.topic.as_deref().unwrap_or(path.as_str());
if ctx
.dedup
.is_duplicate(header.producer, scope, header.sequence)
{
ctx.shared.duplicates.fetch_add(1, Ordering::Relaxed);
tracing::debug!(path, sequence = ?header.sequence, "duplicate suppressed");
return drain(stream).await;
}
let meta = IncomingMeta::from_header(&header, ctx.peer.clone());
if ctx.reorder.enabled() {
let held = Held {
stream,
meta,
path: path.clone(),
};
for (held, gap) in ctx.reorder.admit(scope, header.sequence, held) {
let transfer = IncomingTransfer::new(
held.stream,
Arc::new(held.meta.with_gap(gap)),
Arc::clone(ctx),
);
dispatch(ctx, &held.path, transfer).await;
}
return Ok(());
}
let gap = header.sequence.and_then(|seq| ctx.gaps.observe(scope, seq));
dispatch(
ctx,
&path,
IncomingTransfer::new(stream, Arc::new(meta.with_gap(gap)), Arc::clone(ctx)),
)
.await;
Ok(())
}
pub(crate) struct Held {
stream: RecvHalf,
meta: IncomingMeta,
path: String,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct Refusal {
pub(crate) code: ErrorCode,
pub(crate) stop: u64,
}
impl Refusal {
pub(crate) const UNKNOWN: Refusal = Refusal {
code: ErrorCode::UnknownEndpoint,
stop: codes::UNKNOWN_ENDPOINT,
};
pub(crate) const WRONG_SHAPE: Refusal = Refusal {
code: ErrorCode::Unsupported,
stop: codes::UNSUPPORTED,
};
pub(crate) const PAIR_TAKEN: Refusal = Refusal {
code: ErrorCode::Rejected,
stop: codes::LIMIT_EXCEEDED,
};
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum Wanted {
OneWay,
Exchange,
}
pub(crate) fn refusal_for(route: Option<&Route>, wanted: Wanted) -> Option<Refusal> {
match (route, wanted) {
(None, _) => Some(Refusal::UNKNOWN),
(Some(Route::Raw(_)), _) => None,
(Some(Route::Transfer(_) | Route::Pair { .. }), Wanted::OneWay) => None,
(Some(Route::Request(_)), Wanted::Exchange) => None,
(Some(Route::Request(_) | Route::Transfer(_) | Route::Pub | Route::Pair { .. }), _) => {
Some(Refusal::WRONG_SHAPE)
}
}
}
async fn dispatch(ctx: &ConnHandle, path: &str, transfer: IncomingTransfer) {
let route = ctx.namespace.lookup(path);
if let Some(refusal) = refusal_for(route.as_ref(), Wanted::OneWay) {
tracing::debug!(path, ?refusal, "the path does not serve a one-way transfer");
transfer.refuse(refusal.stop);
return;
}
match route {
Some(Route::Transfer(queue)) => {
if let Err(e) = queue.send(transfer).await {
tracing::debug!(path, "endpoint went away while dispatching");
e.0.refuse(Refusal::UNKNOWN.stop);
}
}
Some(Route::Raw(queue)) => {
if let Err(e) = queue.send(Incoming::Stream(transfer)).await {
tracing::debug!(path, "acceptor went away while dispatching");
if let Incoming::Stream(t) = e.0 {
t.refuse(Refusal::UNKNOWN.stop);
}
}
}
Some(Route::Pair { queue, owner }) => {
if !owner.claim(ctx) {
tracing::debug!(path, "a paired endpoint already has its peer");
transfer.refuse(Refusal::PAIR_TAKEN.stop);
return;
}
if let Err(e) = queue.send(transfer).await {
tracing::debug!(path, "endpoint went away while dispatching");
e.0.refuse(Refusal::UNKNOWN.stop);
}
}
Some(Route::Request(_) | Route::Pub) | None => {
unreachable!("refusal_for refuses every route that cannot serve a one-way transfer")
}
}
}
async fn drain(mut stream: RecvHalf) -> Result<(), Error> {
let mut scratch = vec![0u8; 8 * 1024];
while stream.read(&mut scratch).await?.is_some() {}
Ok(())
}
async fn handle_bi(ctx: &ConnHandle, send: SendHalf, mut recv: RecvHalf) -> Result<(), Error> {
let preamble = match read_preamble(&mut recv, ctx.limits.max_header_bytes).await {
Ok(preamble) => preamble,
Err(Error::Protocol(reason)) => return violation(ctx, &reason),
Err(e) => return Err(e),
};
if preamble.kind != FrameKind::Data {
return violation(
ctx,
&format!("{} may not open a bidirectional stream", preamble.kind),
);
}
ctx.negotiated().await?;
let header = read_header(&mut recv, &preamble).await?;
let header = match DataHeader::decode(&header) {
Ok(h) => h,
Err(e) => return violation(ctx, &e.to_string()),
};
let Some(path) = header.endpoint.clone() else {
return violation(
ctx,
"the initiating half of an exchange must name an endpoint",
);
};
let route = ctx.namespace.lookup(&path);
let request = IncomingRequest::new(
IncomingTransfer::new(
recv,
Arc::new(IncomingMeta::from_header(&header, ctx.peer.clone())),
Arc::clone(ctx),
),
send,
Arc::clone(ctx),
);
if let Some(refusal) = refusal_for(route.as_ref(), Wanted::Exchange) {
tracing::debug!(path, ?refusal, "the path does not serve an exchange");
request.refuse_coded(refusal.code, refusal.stop).await;
return Ok(());
}
match route {
Some(Route::Request(queue)) => {
if queue.send(request).await.is_err() {
tracing::debug!(path, "endpoint went away while dispatching");
}
}
Some(Route::Raw(queue)) => {
if queue.send(Incoming::Exchange(request)).await.is_err() {
tracing::debug!(path, "acceptor went away while dispatching");
}
}
Some(Route::Transfer(_) | Route::Pub | Route::Pair { .. }) | None => {
unreachable!("refusal_for refuses every route that cannot serve an exchange")
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::{Refusal, Wanted, refusal_for};
use crate::listener::{PairOwner, Route};
use std::sync::Arc;
use tokio::sync::mpsc;
fn transfer() -> Route {
Route::Transfer(mpsc::channel(1).0)
}
fn request() -> Route {
Route::Request(mpsc::channel(1).0)
}
fn raw() -> Route {
Route::Raw(mpsc::channel(1).0)
}
fn pair() -> Route {
Route::Pair {
queue: mpsc::channel(1).0,
owner: Arc::new(PairOwner::new()),
}
}
#[test]
fn every_route_mismatch_earns_the_refusal_the_protocol_names() {
assert_eq!(refusal_for(None, Wanted::OneWay), Some(Refusal::UNKNOWN));
assert_eq!(refusal_for(None, Wanted::Exchange), Some(Refusal::UNKNOWN));
assert_eq!(refusal_for(Some(&raw()), Wanted::OneWay), None);
assert_eq!(refusal_for(Some(&raw()), Wanted::Exchange), None);
for route in [transfer(), pair()] {
assert_eq!(refusal_for(Some(&route), Wanted::OneWay), None);
assert_eq!(
refusal_for(Some(&route), Wanted::Exchange),
Some(Refusal::WRONG_SHAPE)
);
}
assert_eq!(refusal_for(Some(&request()), Wanted::Exchange), None);
assert_eq!(
refusal_for(Some(&request()), Wanted::OneWay),
Some(Refusal::WRONG_SHAPE)
);
assert_eq!(
refusal_for(Some(&Route::Pub), Wanted::OneWay),
Some(Refusal::WRONG_SHAPE)
);
assert_eq!(
refusal_for(Some(&Route::Pub), Wanted::Exchange),
Some(Refusal::WRONG_SHAPE)
);
}
#[test]
fn the_three_refusals_are_distinct_on_the_wire() {
use weida_protocol::codes;
assert_eq!(Refusal::UNKNOWN.stop, codes::UNKNOWN_ENDPOINT);
assert_eq!(Refusal::WRONG_SHAPE.stop, codes::UNSUPPORTED);
assert_eq!(Refusal::PAIR_TAKEN.stop, codes::LIMIT_EXCEEDED);
assert_ne!(Refusal::UNKNOWN, Refusal::WRONG_SHAPE);
assert_ne!(Refusal::WRONG_SHAPE, Refusal::PAIR_TAKEN);
}
}