use std::cell::RefCell;
use std::collections::{BTreeMap, VecDeque};
use std::net::SocketAddr;
use std::rc::Rc;
use tokio::sync::{mpsc, oneshot};
use crate::constants;
use crate::core::{
ConnEvent, ConnOutput, Connection as CoreConnection, ConnectionId, Dir, Disposition,
EndpointOutput, IntroId, ToEndpoint, Transmit,
};
use crate::error::{AcceptError, ConnectError, ConnectionLost};
use crate::identity::{Identity, PublicKeyOf};
use crate::packet::Handshake;
use super::connection::Connection;
use super::shared::{
Command, ConnCell, NotificationSlots, PendingOutcome, PendingSlot, Shell, ShellLink, Wakers,
now, resolve_slot, wake_settled,
};
use super::staged::Intro;
use super::wire::Wire;
const DRAIN_BOUND: usize = 100_000;
struct ConnRecord<I: Identity> {
cell: Rc<RefCell<ConnCell<I::Suite>>>,
remote_static: PublicKeyOf<I>,
slot: Option<Rc<RefCell<PendingSlot<I>>>>,
}
#[derive(Debug, Clone, Copy)]
struct ReadyIntro {
id: IntroId,
source: SocketAddr,
}
enum Event<I: Identity> {
Command(Option<Command<I>>),
Received(std::io::Result<(usize, SocketAddr)>),
Timeout,
}
pub(crate) struct Driver<I: Identity, W: Wire> {
wire: W,
shell: Shell<I>,
commands: mpsc::UnboundedReceiver<Command<I>>,
conns: BTreeMap<ConnectionId, ConnRecord<I>>,
ready: VecDeque<ReadyIntro>,
waiting: VecDeque<oneshot::Sender<Intro<I>>>,
link: Rc<dyn ShellLink>,
last_timeout: Option<std::time::Instant>,
overdue_streak: std::cell::Cell<u32>,
}
impl<I: Identity + 'static, W: Wire> Driver<I, W> {
pub(crate) fn new(
wire: W,
shell: Shell<I>,
commands: mpsc::UnboundedReceiver<Command<I>>,
) -> Self {
Self {
wire,
link: Rc::new(shell.clone()),
shell,
commands,
last_timeout: None,
overdue_streak: std::cell::Cell::new(0),
conns: BTreeMap::new(),
ready: VecDeque::new(),
waiting: VecDeque::new(),
}
}
pub(crate) async fn run(mut self) {
let mut buf = [0u8; constants::MAX_DATAGRAM];
loop {
let outgoing = self.serve();
let deadline = self.deadline();
self.transmit(outgoing).await;
let handles = self.shell.state.borrow().handles;
if handles == 0 {
break;
}
let event = {
let Self { wire, commands, .. } = &mut self;
tokio::select! {
biased;
command = commands.recv() => Event::Command(command),
received = wire.recv_from(&mut buf) => Event::Received(received),
() = sleep_until(deadline), if deadline.is_some() => Event::Timeout,
}
};
if !matches!(event, Event::Timeout) {
self.last_timeout = None;
self.overdue_streak.set(0);
}
match event {
Event::Command(Some(command)) => self.handle_command(command),
Event::Command(None) => break,
Event::Received(Ok((len, src))) => {
let len = len.min(buf.len());
self.handle_datagram(src, &buf[..len]);
}
Event::Received(Err(error)) => {
tracing::warn!(
target: "slither::io",
verb = "recv_from",
%error,
"Wire::recv_from failed; the endpoint is untouched",
);
}
Event::Timeout => self.handle_timeout(),
}
}
}
fn serve(&mut self) -> Vec<Outgoing> {
let mut out = Vec::new();
let mut settled = false;
for _ in 0..DRAIN_BOUND {
self.serve_endpoint(&mut out);
let dirty: Vec<ConnectionId> = self
.conns
.iter()
.filter(|(_, record)| record.cell.borrow().dirty)
.map(|(id, _)| *id)
.collect();
if dirty.is_empty() {
settled = true;
break;
}
for id in dirty {
self.serve_connection(id, &mut out);
}
}
assert!(
settled,
"the shell's drain did not settle in {DRAIN_BOUND} passes (§16.4)"
);
self.release_dead();
self.prune_ready();
self.prune_waiting();
self.dispatch_intros();
out
}
fn prune_ready(&mut self) {
if self.ready.is_empty() {
return;
}
let state = self.shell.state.borrow();
self.ready
.retain(|ready| state.endpoint.intro_source(ready.id).is_some());
}
fn prune_waiting(&mut self) {
self.waiting.retain(|reply| !reply.is_closed());
}
fn serve_endpoint(&mut self, out: &mut Vec<Outgoing>) {
for _ in 0..DRAIN_BOUND {
let output = self.shell.state.borrow_mut().endpoint.poll_output();
match output {
EndpointOutput::Timeout(_) => return,
EndpointOutput::Transmit(transmit) => out.push(Outgoing {
conn: None,
transmit,
}),
EndpointOutput::IntroReady(id, source) => {
self.ready.push_back(ReadyIntro { id, source });
}
EndpointOutput::ToConnection(id, install) => {
if let Some(record) = self.conns.get(&id) {
let mut cell = record.cell.borrow_mut();
if let Some(core) = cell.core.as_mut() {
core.handle_endpoint_event(now(), install);
}
cell.dirty = true;
}
}
EndpointOutput::HandshakeFailed(id, error) => self.fail_pending(id, error),
EndpointOutput::Replaced(id) => {
self.deliver_to_core(id, CoreConnection::replaced);
}
EndpointOutput::Contested(id) => {
self.deliver_to_core(id, CoreConnection::mark_contested);
}
}
}
panic!(
"core::Endpoint::poll_output did not reach Timeout in {DRAIN_BOUND} outputs (§16.4)"
);
}
fn deliver_to_core(
&mut self,
id: ConnectionId,
verb: impl FnOnce(&mut CoreConnection<I::Suite>, std::time::Instant),
) {
let Some(record) = self.conns.get(&id) else {
return;
};
let mut cell = record.cell.borrow_mut();
if let Some(core) = cell.core.as_mut() {
verb(core, now());
}
cell.dirty = true;
}
fn serve_connection(&mut self, id: ConnectionId, out: &mut Vec<Outgoing>) {
let Some(record) = self.conns.get(&id) else {
return;
};
let cell = Rc::clone(&record.cell);
for _ in 0..DRAIN_BOUND {
let output = {
let mut borrow = cell.borrow_mut();
borrow.dirty = false;
match borrow.core.as_mut() {
Some(core) => core.poll_output(),
None => return,
}
};
match output {
ConnOutput::Timeout(_) => return,
ConnOutput::Transmit(transmit) => out.push(Outgoing {
conn: Some(id),
transmit,
}),
ConnOutput::ToEndpoint(event) => {
debug_assert!(
matches!(event, ToEndpoint::Retired { .. }),
"§16.4 defines exactly one connection→endpoint event",
);
self.shell
.state
.borrow_mut()
.endpoint
.handle_connection_event(now(), id, event);
}
ConnOutput::Event(event) => self.publish(id, &cell, event),
}
}
panic!(
"core::Connection::poll_output did not reach Timeout in {DRAIN_BOUND} outputs (§16.4)"
);
}
fn publish(
&mut self,
id: ConnectionId,
cell: &Rc<RefCell<ConnCell<I::Suite>>>,
event: ConnEvent,
) {
match event {
ConnEvent::Established => self.establish(id, cell),
ConnEvent::Closed(lost) => Self::latch(cell, lost),
ConnEvent::StreamReadable { r } => {
Self::wake_stream(cell, |cell| cell.blocked_readers.get_mut(&r));
}
ConnEvent::StreamWritable { r } => {
Self::wake_stream(cell, |cell| cell.blocked_writers.get_mut(&r));
}
ConnEvent::StreamReset { r, error_code } => {
let gates_our_send_half = cell
.borrow()
.core
.as_ref()
.and_then(|core| core.stream_id(r))
.is_some_and(|id| id.dir() == Dir::Uni);
if gates_our_send_half {
cell.borrow_mut().peer_resets.insert(r, error_code);
}
Self::wake_stream(cell, |cell| cell.blocked_readers.get_mut(&r));
Self::wake_stream(cell, |cell| cell.blocked_writers.get_mut(&r));
Self::wake_stream(cell, |cell| cell.blocked_ackers.get_mut(&r));
wake_settled(cell);
}
ConnEvent::StreamOpened { dir } => {
Self::wake_stream(cell, |cell| Some(&mut cell.stream_acceptors[dir.slot()]));
}
ConnEvent::StreamsAvailable { dir } => {
Self::wake_stream(cell, |cell| Some(&mut cell.stream_openers[dir.slot()]));
if dir == Dir::Uni {
Self::wake_stream(cell, |cell| Some(&mut cell.message_senders));
}
}
ConnEvent::SendCreditAvailable => {
Self::wake_stream(cell, |cell| Some(&mut cell.message_senders));
}
ConnEvent::MessageReadable => {
Self::wake_stream(cell, |cell| Some(&mut cell.message_readers));
}
ConnEvent::DatagramReadable => {
Self::wake_stream(cell, |cell| Some(&mut cell.datagram_readers));
}
ConnEvent::StreamFinished { r } => {
let woken = {
let mut borrow = cell.borrow_mut();
let mut woken = borrow.note_send_finished(r);
woken.extend(borrow.settled_wakers.take_all());
woken
};
for waker in woken {
waker.wake();
}
}
ConnEvent::AddressMoved { from, to } => {
Self::notify(cell, |slots| slots.address_moved(from, to), Some(to));
}
ConnEvent::Contested => {
Self::notify(cell, NotificationSlots::contested, None);
}
ConnEvent::ContestCleared => {
Self::notify(cell, NotificationSlots::contest_cleared, None);
}
}
}
fn notify(
cell: &Rc<RefCell<ConnCell<I::Suite>>>,
fill: impl FnOnce(&mut NotificationSlots),
anchor: Option<SocketAddr>,
) {
let woken = {
let mut borrow = cell.borrow_mut();
if let Some(anchor) = anchor {
borrow.remote_address = anchor;
}
fill(&mut borrow.notifications);
borrow.notification_wakers.take_all()
};
for waker in woken {
waker.wake();
}
}
fn wake_stream(
cell: &Rc<RefCell<ConnCell<I::Suite>>>,
select: impl FnOnce(&mut ConnCell<I::Suite>) -> Option<&mut Wakers>,
) {
let woken = {
let mut borrow = cell.borrow_mut();
select(&mut borrow)
.map(Wakers::take_all)
.unwrap_or_default()
};
for waker in woken {
waker.wake();
}
}
fn establish(&mut self, id: ConnectionId, cell: &Rc<RefCell<ConnCell<I::Suite>>>) {
let Some(record) = self.conns.get_mut(&id) else {
return;
};
let anchor = {
let borrow = cell.borrow();
match borrow.core.as_ref().and_then(CoreConnection::session) {
Some(session) => session.anchor,
None => {
debug_assert!(false, "ConnEvent::Established without an installed session");
return;
}
}
};
cell.borrow_mut().remote_address = anchor;
let Some(slot) = record.slot.take() else {
return;
};
let session_id = match session_id_of(cell) {
Some(session_id) => session_id,
None => return,
};
let handle = Connection::new(
Rc::clone(&self.link),
Rc::clone(cell),
id,
record.remote_static.clone(),
session_id,
);
resolve_slot(&slot, PendingOutcome::Ready(handle));
}
fn fail_pending(&mut self, id: ConnectionId, error: ConnectError) {
let Some(record) = self.conns.remove(&id) else {
return;
};
if let Some(slot) = record.slot {
resolve_slot(&slot, PendingOutcome::Failed(error));
}
}
fn release_dead(&mut self) {
let done: Vec<ConnectionId> = self
.conns
.iter()
.filter(|(_, record)| {
let cell = record.cell.borrow();
cell.closed.is_some() && !cell.is_established()
})
.map(|(id, _)| *id)
.collect();
for id in done {
let Some(record) = self.conns.remove(&id) else {
continue;
};
record.cell.borrow_mut().core = None;
if let Some(slot) = record.slot {
resolve_slot(&slot, PendingOutcome::Failed(ConnectError::Local));
}
}
}
fn dispatch_intros(&mut self) {
while !self.ready.is_empty() {
while self.waiting.front().is_some_and(oneshot::Sender::is_closed) {
self.waiting.pop_front();
}
let Some(reply) = self.waiting.pop_front() else {
return;
};
let ready = self
.ready
.pop_front()
.expect("the loop condition checked it");
let sender_index = self
.shell
.state
.borrow()
.endpoint
.intro_sender_index(ready.id)
.unwrap_or_default();
let intro = Intro::new(self.shell.clone(), ready.id, ready.source, sender_index);
drop(reply.send(intro));
}
}
async fn transmit(&mut self, outgoing: Vec<Outgoing>) {
for Outgoing { conn, transmit } in outgoing {
let Transmit { to, data } = transmit;
if let Err(error) = self.wire.send_to(&data, to).await {
tracing::warn!(
target: "slither::io",
verb = "send_to",
conn = ?conn,
to = %to,
%error,
"Wire::send_to failed; the connection is untouched (§16.3, ruling 49)",
);
}
}
}
fn handle_datagram(&mut self, src: SocketAddr, datagram: &[u8]) {
let now = now();
let disposition = self
.shell
.state
.borrow_mut()
.endpoint
.handle_datagram(now, src, datagram);
if let Disposition::ForConnection(id) = disposition
&& let Some(record) = self.conns.get(&id)
{
let cell = Rc::clone(&record.cell);
{
let mut borrow = cell.borrow_mut();
if let Some(core) = borrow.core.as_mut() {
core.handle_datagram(now, src, datagram);
}
borrow.dirty = true;
}
wake_settled(&cell);
}
}
fn handle_timeout(&mut self) {
let now = now();
self.last_timeout = Some(now);
self.shell.state.borrow_mut().endpoint.handle_timeout(now);
for record in self.conns.values() {
let mut cell = record.cell.borrow_mut();
if let Some(core) = cell.core.as_mut() {
core.handle_timeout(now);
}
cell.dirty = true;
}
}
fn deadline(&self) -> Option<std::time::Instant> {
let endpoint = self.shell.state.borrow().endpoint.next_deadline();
self.conns
.values()
.filter_map(|record| record.cell.borrow().core.as_ref()?.next_deadline())
.chain(endpoint)
.min()
.inspect(|announced| {
if let Some(fired_at) = self.last_timeout {
if *announced <= fired_at {
let streak = self.overdue_streak.get() + 1;
self.overdue_streak.set(streak);
assert!(streak < 3, "{PAST_DEADLINE}");
} else {
self.overdue_streak.set(0);
}
}
})
}
fn handle_command(&mut self, command: Command<I>) {
match command {
Command::Connect {
id,
core,
remote,
remote_static,
slot,
} => self.command_connect(id, *core, remote, remote_static, slot),
Command::Cancel(id) => self.command_cancel(id),
Command::Accept(reply) => {
self.waiting.push_back(reply);
self.dispatch_intros();
}
Command::ReadIdentity(id, reply) => {
let result = self
.shell
.state
.borrow_mut()
.endpoint
.read_identity(now(), id);
drop(reply.send(result));
}
Command::Authenticate(id, psk, reply) => {
let result = self
.shell
.state
.borrow_mut()
.endpoint
.authenticate(now(), id, &psk);
drop(reply.send(result));
}
Command::AcceptChain(id, remote_static, reply) => {
self.command_accept_chain(id, remote_static, reply);
}
Command::Reject(id) => self.shell.state.borrow_mut().endpoint.reject(now(), id),
Command::Dirty(id) => {
if let Some(record) = self.conns.get(&id) {
record.cell.borrow_mut().dirty = true;
}
}
Command::HandlesGone => {}
}
}
fn command_connect(
&mut self,
id: ConnectionId,
core: CoreConnection<I::Suite>,
remote: SocketAddr,
remote_static: PublicKeyOf<I>,
slot: Rc<RefCell<PendingSlot<I>>>,
) {
let cell = Rc::new(RefCell::new(ConnCell::new(core, remote)));
self.conns.insert(
id,
ConnRecord {
cell,
remote_static,
slot: Some(slot),
},
);
self.shell
.state
.borrow_mut()
.endpoint
.start_attempt(now(), id);
}
fn command_cancel(&mut self, id: ConnectionId) {
self.conns.remove(&id);
}
fn command_accept_chain(
&mut self,
id: IntroId,
remote_static: PublicKeyOf<I>,
reply: oneshot::Sender<Result<Connection<I::Suite>, AcceptError>>,
) {
let accepted = self.shell.state.borrow_mut().endpoint.accept(now(), id);
let (conn_id, core) = match accepted {
Ok(accepted) => accepted,
Err(error) => {
drop(reply.send(Err(error)));
return;
}
};
let established = core.session().map(|session| {
(
session.anchor,
<I::Suite as Handshake>::session_id(&session.seal).clone(),
)
});
let Some((anchor, session_id)) = established else {
debug_assert!(false, "§16.4: `accept()` returns an established connection");
drop(reply.send(Err(AcceptError::Stale)));
return;
};
let cell = Rc::new(RefCell::new(ConnCell::new(core, anchor)));
let handle = Connection::new(
Rc::clone(&self.link),
Rc::clone(&cell),
conn_id,
remote_static.clone(),
session_id,
);
self.conns.insert(
conn_id,
ConnRecord {
cell,
remote_static,
slot: None,
},
);
drop(reply.send(Ok(handle)));
}
}
impl<I: Identity, W: Wire> Driver<I, W> {
fn latch(cell: &Rc<RefCell<ConnCell<I::Suite>>>, lost: ConnectionLost) {
let woken = {
let mut borrow = cell.borrow_mut();
if borrow.closed.is_some() {
return;
}
borrow.closed = Some(lost);
let mut woken = borrow.closed_wakers.take_all();
woken.extend(borrow.notification_wakers.take_all());
woken.extend(borrow.take_all_stream_wakers());
woken
};
for waker in woken {
waker.wake();
}
}
fn stop(&mut self) {
self.shell.state.borrow_mut().driver_stopped = true;
self.commands.close();
while let Ok(command) = self.commands.try_recv() {
if let Command::Connect { slot, .. } = command {
resolve_slot(&slot, PendingOutcome::Failed(ConnectError::Local));
}
}
for record in self.conns.values_mut() {
Self::latch(&record.cell, ConnectionLost::EndpointDropped);
record.cell.borrow_mut().core = None;
if let Some(slot) = record.slot.take() {
resolve_slot(&slot, PendingOutcome::Failed(ConnectError::Local));
}
}
self.waiting.clear();
self.conns.clear();
}
}
impl<I: Identity, W: Wire> Drop for Driver<I, W> {
fn drop(&mut self) {
self.stop();
}
}
struct Outgoing {
conn: Option<ConnectionId>,
transmit: Transmit,
}
fn session_id_of<S: Handshake>(cell: &Rc<RefCell<ConnCell<S>>>) -> Option<hiss::noise::SessionId> {
let borrow = cell.borrow();
let session = borrow.core.as_ref()?.session()?;
Some(<S as Handshake>::session_id(&session.seal).clone())
}
const PAST_DEADLINE: &str = "a core announced a deadline in the past — ruling 141's spin class: `sleep_until` \
completes at once, `handle_timeout` re-fires, and the actor spins";
async fn sleep_until(deadline: Option<std::time::Instant>) {
match deadline {
Some(deadline) => tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)).await,
None => std::future::pending().await,
}
}