use std::cell::RefCell;
use std::future::{Future, poll_fn};
use std::net::SocketAddr;
use std::pin::Pin;
use std::rc::Rc;
use std::task::{Context, Poll};
use crate::config::Config;
use crate::core::{ConnectionId, Endpoint as CoreEndpoint, ToEndpoint};
use crate::error::ConnectError;
use crate::identity::{Identity, PublicKeyOf};
use super::connection::Connection;
use super::driver::Driver;
use super::shared::{Command, PendingOutcome, PendingSlot, Shell, now};
use super::staged::Intro;
use super::wire::Wire;
pub struct Endpoint<I: Identity> {
shell: Shell<I>,
}
impl<I: Identity> Endpoint<I> {
pub fn builder<W: Wire>() -> EndpointBuilder<I, W> {
EndpointBuilder::new()
}
pub async fn accept(&self) -> Option<Intro<I>> {
let mut pending = None;
poll_fn(|cx| self.poll_accept(cx, &mut pending)).await
}
pub(crate) fn poll_accept(
&self,
cx: &mut Context<'_>,
pending: &mut Option<tokio::sync::oneshot::Receiver<Intro<I>>>,
) -> Poll<Option<Intro<I>>> {
let rx = match pending {
Some(rx) => rx,
None => {
if self.shell.driver_stopped() {
return Poll::Ready(None);
}
let (tx, rx) = tokio::sync::oneshot::channel();
self.shell.send(Command::Accept(tx));
pending.insert(rx)
}
};
match Pin::new(rx).poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(outcome) => {
*pending = None;
Poll::Ready(outcome.ok())
}
}
}
}
impl<I: Identity> Endpoint<I>
where
I::Suite: crate::packet::Handshake<Psk = ()>,
{
pub fn connect(
&self,
remote: SocketAddr,
remote_static: PublicKeyOf<I>,
) -> Result<Connecting<I>, ConnectError> {
self.connect_with(remote, remote_static, ())
}
}
impl<I: Identity> Endpoint<I> {
pub fn connect_with(
&self,
remote: SocketAddr,
remote_static: PublicKeyOf<I>,
psk: crate::identity::PskOf<I>,
) -> Result<Connecting<I>, ConnectError> {
let (id, core) = {
let mut state = self.shell.state.borrow_mut();
if state.driver_stopped {
return Err(ConnectError::Local);
}
let minted = state
.endpoint
.mint_pending(now(), remote, remote_static.clone(), psk)?;
state.drain_endpoint();
minted
};
let slot = Rc::new(RefCell::new(PendingSlot::new()));
self.shell.send(Command::Connect {
id,
core: Box::new(core),
remote,
remote_static,
slot: Rc::clone(&slot),
});
Ok(Connecting::new(self.shell.clone(), slot, id))
}
}
impl<I: Identity> std::fmt::Debug for Endpoint<I> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Endpoint").finish_non_exhaustive()
}
}
impl<I: Identity> Drop for Endpoint<I> {
fn drop(&mut self) {
self.shell.release();
}
}
pub struct Connecting<I: Identity> {
shell: Shell<I>,
slot: Rc<RefCell<PendingSlot<I>>>,
id: ConnectionId,
resolved: bool,
}
impl<I: Identity> Connecting<I> {
fn new(shell: Shell<I>, slot: Rc<RefCell<PendingSlot<I>>>, id: ConnectionId) -> Self {
shell.acquire();
Self {
shell,
slot,
id,
resolved: false,
}
}
}
impl<I: Identity> Future for Connecting<I> {
type Output = Result<Connection<I::Suite>, ConnectError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
let mut slot = this.slot.borrow_mut();
match std::mem::replace(&mut slot.outcome, PendingOutcome::Waiting) {
PendingOutcome::Waiting => {
slot.waker = Some(cx.waker().clone());
Poll::Pending
}
PendingOutcome::Ready(connection) => {
this.resolved = true;
Poll::Ready(Ok(connection))
}
PendingOutcome::Failed(error) => {
this.resolved = true;
Poll::Ready(Err(error))
}
}
}
}
impl<I: Identity> std::fmt::Debug for Connecting<I> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Connecting").finish_non_exhaustive()
}
}
impl<I: Identity> Drop for Connecting<I> {
fn drop(&mut self) {
let in_flight =
!self.resolved && matches!(self.slot.borrow().outcome, PendingOutcome::Waiting);
if in_flight {
{
let mut state = self.shell.state.borrow_mut();
state.endpoint.handle_connection_event(
now(),
self.id,
ToEndpoint::Retired { our_index: 0 },
);
state.drain_endpoint();
}
self.shell.send(Command::Cancel(self.id));
}
self.shell.release();
}
}
pub struct EndpointBuilder<I: Identity, W: Wire> {
identity: Option<I>,
wire: Option<W>,
config: Config,
rng_seed: Option<[u8; 32]>,
}
impl<I: Identity, W: Wire> EndpointBuilder<I, W> {
fn new() -> Self {
Self {
identity: None,
wire: None,
config: Config::new(),
rng_seed: None,
}
}
#[must_use]
pub fn identity(mut self, identity: I) -> Self {
self.identity = Some(identity);
self
}
#[must_use]
pub fn wire(mut self, wire: W) -> Self {
self.wire = Some(wire);
self
}
#[must_use]
pub fn config(mut self, config: Config) -> Self {
self.config = config;
self
}
#[must_use]
pub fn rng_seed(mut self, seed: [u8; 32]) -> Self {
self.rng_seed = Some(seed);
self
}
#[must_use]
pub fn build(self) -> Endpoint<I>
where
I: 'static,
W: 'static,
{
let identity = self
.identity
.expect("Endpoint::builder() requires an identity");
let wire = self.wire.expect("Endpoint::builder() requires a wire");
let rng_seed = self.rng_seed.unwrap_or_else(|| {
let mut seed = [0u8; 32];
getrandom::fill(&mut seed).expect("OS entropy for the endpoint RNG (§16.6)");
seed
});
let core = CoreEndpoint::new(super::shared::now(), self.config, identity, rng_seed);
let (shell, commands) = Shell::new(core);
let endpoint = Endpoint {
shell: shell.clone(),
};
endpoint.shell.acquire();
let driver = Driver::new(wire, shell, commands).run();
let spawned = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
tokio::task::spawn_local(driver);
}));
if spawned.is_err() {
panic!(
"slither: Endpoint::builder()…build() was called outside a \
`tokio::task::LocalSet`. The driver is a single `!Send` task \
(a DH provider and a Wire are not required to be `Send`), so \
it is spawned with `tokio::task::spawn_local`, which needs a \
current-thread runtime with a `LocalSet` — `#[tokio::main]` \
alone is not one. Wrap your code in \
`slither::prelude::block_on(async {{ … }})`, or build the \
endpoint inside \
`tokio::task::LocalSet::new().run_until(…)`."
);
}
endpoint
}
}
impl<I: Identity, W: Wire> Default for EndpointBuilder<I, W> {
fn default() -> Self {
Self::new()
}
}
impl<I: Identity, W: Wire> std::fmt::Debug for EndpointBuilder<I, W> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("EndpointBuilder")
.field("identity", &self.identity.is_some())
.field("wire", &self.wire.is_some())
.field("seeded", &self.rng_seed.is_some())
.finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use rand_chacha::ChaCha20Rng;
use rand_chacha::rand_core::SeedableRng;
use super::Endpoint;
use crate::identity::SoftwareIdentity;
use crate::packet::ReferenceSuite;
#[test]
fn build_outside_a_localset_panics_with_slithers_own_message() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("a current-thread runtime");
let payload = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
rt.block_on(async {
let identity: SoftwareIdentity<ReferenceSuite> =
SoftwareIdentity::generate(ChaCha20Rng::from_seed([0x44; 32]))
.expect("generate an identity");
let socket = tokio::net::UdpSocket::bind("127.0.0.1:0")
.await
.expect("bind a socket");
let _endpoint = Endpoint::builder().identity(identity).wire(socket).build();
});
}))
.expect_err("build() outside a LocalSet must panic");
let message = payload
.downcast_ref::<String>()
.map(String::as_str)
.or_else(|| payload.downcast_ref::<&'static str>().copied())
.expect("a string panic payload");
assert!(
message.contains("slither::prelude::block_on"),
"the guard must name the fix; got: {message}"
);
assert!(
message.contains("tokio::task::LocalSet"),
"the guard must name what is missing; got: {message}"
);
assert!(
message.contains("Endpoint::builder"),
"the guard must name the call site; got: {message}"
);
}
}