use std::cell::RefCell;
use std::future::poll_fn;
use std::net::SocketAddr;
use std::rc::Rc;
use std::task::{Context, Poll};
use std::time::Duration;
use crate::constants;
use crate::core::connection::{AckSnapshot, SendMessage, validate_persistent_keepalive};
use crate::core::{Connection as CoreConnection, ConnectionId, Dir, StreamRef};
use crate::error::{ConfigError, ConnectionLost, DatagramError, MessageError};
use crate::packet::{Channel, Handshake};
use super::shared::{ConnCell, Notification, ShellLink, WakerSlot, close_now, now};
use super::stream::{BiStream, RecvStream, SendStream};
type PublicKeyFor<S> = <<S as Channel>::Curve as hiss::curve::Curve>::PublicKey;
pub struct Connection<S: Handshake> {
shell: Rc<dyn ShellLink>,
cell: Rc<RefCell<ConnCell<S>>>,
id: ConnectionId,
remote_static: PublicKeyFor<S>,
session_id: hiss::noise::SessionId,
}
impl<S: Handshake> Connection<S> {
pub(crate) fn new(
shell: Rc<dyn ShellLink>,
cell: Rc<RefCell<ConnCell<S>>>,
id: ConnectionId,
remote_static: PublicKeyFor<S>,
session_id: hiss::noise::SessionId,
) -> Self {
shell.acquire();
cell.borrow_mut().handles += 1;
Self {
shell,
cell,
id,
remote_static,
session_id,
}
}
#[cfg(feature = "tower")]
pub(crate) fn clone_handle(&self) -> Self {
Self::new(
Rc::clone(&self.shell),
Rc::clone(&self.cell),
self.id,
self.remote_static.clone(),
self.session_id.clone(),
)
}
pub async fn close(&self, code: u64, reason: &[u8]) {
poll_fn(|cx| self.poll_close(cx, code, reason)).await
}
fn poll_close(&self, _cx: &mut Context<'_>, code: u64, reason: &[u8]) -> Poll<()> {
self.close_now(code, reason);
Poll::Ready(())
}
fn close_now(&self, code: u64, reason: &[u8]) {
close_now(&self.shell, &self.cell, self.id, code, reason);
}
pub async fn closed(&self) -> ConnectionLost {
let cell = Rc::clone(&self.cell);
let key = cell.borrow_mut().closed_wakers.key();
let slot = WakerSlot::new(key, {
let cell = Rc::clone(&self.cell);
move |key| cell.borrow_mut().closed_wakers.unpark(key)
});
poll_fn(|cx| self.poll_closed(cx, slot.key())).await
}
fn poll_closed(&self, cx: &mut Context<'_>, key: u64) -> Poll<ConnectionLost> {
let mut cell = self.cell.borrow_mut();
match cell.closed.clone() {
Some(lost) => Poll::Ready(lost),
None => {
cell.closed_wakers.park(key, cx);
Poll::Pending
}
}
}
pub async fn notified(&self) -> Result<Notification, ConnectionLost> {
let slot = self.notification_slot();
poll_fn(|cx| self.poll_notified(cx, slot.key())).await
}
pub(crate) fn poll_notified(
&self,
cx: &mut Context<'_>,
key: u64,
) -> Poll<Result<Notification, ConnectionLost>> {
let mut cell = self.cell.borrow_mut();
if let Some(notification) = cell.notifications.take_oldest() {
return Poll::Ready(Ok(notification));
}
if let Some(lost) = cell.closed.clone() {
return Poll::Ready(Err(lost));
}
if cell.core.is_none() {
debug_assert!(
false,
"a connection cell held neither a core nor a close reason (§16.3)"
);
return Poll::Ready(Err(ConnectionLost::EndpointDropped));
}
cell.notification_wakers.park(key, cx);
Poll::Pending
}
fn notification_slot(&self) -> WakerSlot<impl FnMut(u64)> {
let key = self.cell.borrow_mut().notification_wakers.key();
WakerSlot::new(key, {
let cell = Rc::clone(&self.cell);
move |key| cell.borrow_mut().notification_wakers.unpark(key)
})
}
pub fn set_persistent_keepalive(&self, interval: Option<Duration>) -> Result<(), ConfigError> {
validate_persistent_keepalive(interval)?;
let applied = {
let mut cell = self.cell.borrow_mut();
match cell.core.as_mut() {
Some(core) => {
let result = core.set_persistent_keepalive(now(), interval);
debug_assert!(result.is_ok(), "the band was checked immediately above");
cell.dirty = true;
true
}
None => false,
}
};
if applied {
self.shell.mark_dirty(self.id);
}
Ok(())
}
pub fn persistent_keepalive(&self) -> Option<Duration> {
let cell = self.cell.borrow();
cell.core
.as_ref()
.and_then(|core| core.persistent_keepalive())
}
pub async fn acked(&self) -> Result<(), ConnectionLost> {
let snapshot = self
.cell
.borrow()
.core
.as_ref()
.map(CoreConnection::ack_snapshot);
let slot = self.settled_slot();
poll_fn(|cx| self.poll_acked(cx, snapshot.as_ref(), slot.key())).await
}
fn poll_acked(
&self,
cx: &mut Context<'_>,
snapshot: Option<&AckSnapshot>,
key: u64,
) -> Poll<Result<(), ConnectionLost>> {
let mut cell = self.cell.borrow_mut();
if let Some(snap) = snapshot
&& cell
.core
.as_ref()
.is_some_and(|core| core.snapshot_settled(snap))
{
return Poll::Ready(Ok(()));
}
if let Some(lost) = cell.closed.clone() {
return Poll::Ready(Err(lost));
}
if cell.core.is_none() {
debug_assert!(
false,
"a connection cell held neither a core nor a close reason (§16.3)"
);
return Poll::Ready(Err(ConnectionLost::EndpointDropped));
}
cell.settled_wakers.park(key, cx);
Poll::Pending
}
fn settled_slot(&self) -> WakerSlot<impl FnMut(u64)> {
let key = self.cell.borrow_mut().settled_wakers.key();
WakerSlot::new(key, {
let cell = Rc::clone(&self.cell);
move |key| cell.borrow_mut().settled_wakers.unpark(key)
})
}
pub async fn open_bi(&self) -> Result<BiStream<S>, ConnectionLost> {
let slot = self.opener_slot(Dir::Bi);
poll_fn(|cx| self.poll_open_bi(cx, slot.key())).await
}
pub async fn open_uni(&self) -> Result<SendStream<S>, ConnectionLost> {
let slot = self.opener_slot(Dir::Uni);
poll_fn(|cx| self.poll_open_uni(cx, slot.key())).await
}
pub async fn accept_bi(&self) -> Result<BiStream<S>, ConnectionLost> {
let slot = self.acceptor_slot(Dir::Bi);
poll_fn(|cx| self.poll_accept_bi(cx, slot.key())).await
}
pub async fn accept_uni(&self) -> Result<RecvStream<S>, ConnectionLost> {
let slot = self.acceptor_slot(Dir::Uni);
poll_fn(|cx| self.poll_accept_uni(cx, slot.key())).await
}
pub(crate) fn poll_open_bi(
&self,
cx: &mut Context<'_>,
key: u64,
) -> Poll<Result<BiStream<S>, ConnectionLost>> {
self.poll_open_with(cx, key, Dir::Bi, install_bi)
}
pub(crate) fn poll_open_uni(
&self,
cx: &mut Context<'_>,
key: u64,
) -> Poll<Result<SendStream<S>, ConnectionLost>> {
self.poll_open_with(cx, key, Dir::Uni, SendStream::install)
}
pub(crate) fn poll_accept_bi(
&self,
cx: &mut Context<'_>,
key: u64,
) -> Poll<Result<BiStream<S>, ConnectionLost>> {
self.poll_accept_with(cx, key, Dir::Bi, install_bi)
}
pub(crate) fn poll_accept_uni(
&self,
cx: &mut Context<'_>,
key: u64,
) -> Poll<Result<RecvStream<S>, ConnectionLost>> {
self.poll_accept_with(cx, key, Dir::Uni, RecvStream::install)
}
pub async fn send_message(&self, msg: &[u8]) -> Result<(), MessageError> {
let slot = self.message_sender_slot();
poll_fn(|cx| self.poll_send_message(cx, msg, slot.key())).await
}
pub async fn recv_message(&self) -> Result<Vec<u8>, ConnectionLost> {
let slot = self.message_reader_slot();
poll_fn(|cx| self.poll_recv_message(cx, slot.key())).await
}
pub fn send_datagram(&self, data: &[u8]) -> Result<(), DatagramError> {
if data.len() > constants::MAX_DATAGRAM_PAYLOAD {
return Err(DatagramError::TooLarge);
}
let sent = {
let mut cell = self.cell.borrow_mut();
if let Some(lost) = cell.closed.clone() {
return Err(DatagramError::ConnectionLost(lost));
}
let Some(core) = cell.core.as_mut() else {
debug_assert!(
false,
"a connection cell held neither a core nor a close reason (§16.3)"
);
return Err(DatagramError::ConnectionLost(
ConnectionLost::EndpointDropped,
));
};
let sent = core.send_datagram(now(), data);
if sent.is_ok() {
cell.dirty = true;
}
sent
};
if sent.is_ok() {
self.shell.mark_dirty(self.id);
}
sent
}
pub async fn recv_datagram(&self) -> Result<Vec<u8>, ConnectionLost> {
let slot = self.datagram_reader_slot();
poll_fn(|cx| self.poll_recv_datagram(cx, slot.key())).await
}
pub(crate) fn poll_send_message(
&self,
cx: &mut Context<'_>,
msg: &[u8],
key: u64,
) -> Poll<Result<(), MessageError>> {
if msg.len() as u64 > constants::MESSAGE_RECV_MAX {
return Poll::Ready(Err(MessageError::TooLarge));
}
let admitted = {
let mut cell = self.cell.borrow_mut();
if let Some(lost) = cell.closed.clone() {
return Poll::Ready(Err(MessageError::ConnectionLost(lost)));
}
let Some(core) = cell.core.as_mut() else {
debug_assert!(
false,
"a connection cell held neither a core nor a close reason (§16.3)"
);
return Poll::Ready(Err(MessageError::ConnectionLost(
ConnectionLost::EndpointDropped,
)));
};
match core.send_message(now(), msg) {
Err(error) => return Poll::Ready(Err(error)),
Ok(SendMessage::Blocked) => {
cell.message_senders.park(key, cx);
false
}
Ok(SendMessage::Sent) => {
cell.dirty = true;
true
}
}
};
if !admitted {
return Poll::Pending;
}
self.shell.mark_dirty(self.id);
Poll::Ready(Ok(()))
}
pub(crate) fn poll_recv_message(
&self,
cx: &mut Context<'_>,
key: u64,
) -> Poll<Result<Vec<u8>, ConnectionLost>> {
let (outcome, called) = {
let mut cell = self.cell.borrow_mut();
let claimed = cell.core.as_mut().map(|core| core.recv_message(now()));
let called = claimed.is_some();
if called {
cell.dirty = true;
}
let outcome = match claimed {
Some(Some(payload)) => Some(Ok(payload)),
_ => match cell.closed.clone() {
Some(lost) => Some(Err(lost)),
None if cell.core.is_none() => {
debug_assert!(
false,
"a connection cell held neither a core nor a close reason (§16.3)"
);
Some(Err(ConnectionLost::EndpointDropped))
}
None => {
cell.message_readers.park(key, cx);
None
}
},
};
(outcome, called)
};
if called {
self.shell.mark_dirty(self.id);
}
match outcome {
Some(result) => Poll::Ready(result),
None => Poll::Pending,
}
}
pub(crate) fn poll_recv_datagram(
&self,
cx: &mut Context<'_>,
key: u64,
) -> Poll<Result<Vec<u8>, ConnectionLost>> {
let mut cell = self.cell.borrow_mut();
if let Some(payload) = cell.core.as_mut().and_then(CoreConnection::recv_datagram) {
return Poll::Ready(Ok(payload));
}
if let Some(lost) = cell.closed.clone() {
return Poll::Ready(Err(lost));
}
if cell.core.is_none() {
debug_assert!(
false,
"a connection cell held neither a core nor a close reason (§16.3)"
);
return Poll::Ready(Err(ConnectionLost::EndpointDropped));
}
cell.datagram_readers.park(key, cx);
Poll::Pending
}
fn message_reader_slot(&self) -> WakerSlot<impl FnMut(u64)> {
let key = self.cell.borrow_mut().message_readers.key();
WakerSlot::new(key, {
let cell = Rc::clone(&self.cell);
move |key| cell.borrow_mut().message_readers.unpark(key)
})
}
fn datagram_reader_slot(&self) -> WakerSlot<impl FnMut(u64)> {
let key = self.cell.borrow_mut().datagram_readers.key();
WakerSlot::new(key, {
let cell = Rc::clone(&self.cell);
move |key| cell.borrow_mut().datagram_readers.unpark(key)
})
}
fn message_sender_slot(&self) -> WakerSlot<impl FnMut(u64)> {
let key = self.cell.borrow_mut().message_senders.key();
WakerSlot::new(key, {
let cell = Rc::clone(&self.cell);
move |key| cell.borrow_mut().message_senders.unpark(key)
})
}
fn opener_slot(&self, dir: Dir) -> WakerSlot<impl FnMut(u64)> {
let key = self.cell.borrow_mut().stream_openers[dir.slot()].key();
WakerSlot::new(key, {
let cell = Rc::clone(&self.cell);
move |key| cell.borrow_mut().stream_openers[dir.slot()].unpark(key)
})
}
fn acceptor_slot(&self, dir: Dir) -> WakerSlot<impl FnMut(u64)> {
let key = self.cell.borrow_mut().stream_acceptors[dir.slot()].key();
WakerSlot::new(key, {
let cell = Rc::clone(&self.cell);
move |key| cell.borrow_mut().stream_acceptors[dir.slot()].unpark(key)
})
}
fn poll_open_with<T>(
&self,
cx: &mut Context<'_>,
key: u64,
dir: Dir,
build: impl FnOnce(
Rc<dyn ShellLink>,
Rc<RefCell<ConnCell<S>>>,
&mut ConnCell<S>,
ConnectionId,
StreamRef,
) -> T,
) -> Poll<Result<T, ConnectionLost>> {
let mut cell = self.cell.borrow_mut();
if let Some(lost) = cell.closed.clone() {
return Poll::Ready(Err(lost));
}
let Some(core) = cell.core.as_mut() else {
debug_assert!(
false,
"a connection cell held neither a core nor a close reason (§16.3)"
);
return Poll::Ready(Err(ConnectionLost::EndpointDropped));
};
match core.open(dir) {
Ok(r) => {
let handle = build(
Rc::clone(&self.shell),
Rc::clone(&self.cell),
&mut cell,
self.id,
r,
);
Poll::Ready(Ok(handle))
}
Err(_) => {
cell.stream_openers[dir.slot()].park(key, cx);
Poll::Pending
}
}
}
fn poll_accept_with<T>(
&self,
cx: &mut Context<'_>,
key: u64,
dir: Dir,
build: impl FnOnce(
Rc<dyn ShellLink>,
Rc<RefCell<ConnCell<S>>>,
&mut ConnCell<S>,
ConnectionId,
StreamRef,
) -> T,
) -> Poll<Result<T, ConnectionLost>> {
let mut cell = self.cell.borrow_mut();
let claimed = cell.core.as_mut().and_then(|core| core.accept(dir));
if let Some(r) = claimed {
let handle = build(
Rc::clone(&self.shell),
Rc::clone(&self.cell),
&mut cell,
self.id,
r,
);
return Poll::Ready(Ok(handle));
}
if let Some(lost) = cell.closed.clone() {
return Poll::Ready(Err(lost));
}
if cell.core.is_none() {
debug_assert!(
false,
"a connection cell held neither a core nor a close reason (§16.3)"
);
return Poll::Ready(Err(ConnectionLost::EndpointDropped));
}
cell.stream_acceptors[dir.slot()].park(key, cx);
Poll::Pending
}
pub fn remote_static(&self) -> PublicKeyFor<S> {
self.remote_static.clone()
}
pub fn remote_address(&self) -> SocketAddr {
self.cell.borrow().remote_address
}
pub fn session_id(&self) -> hiss::noise::SessionId {
self.session_id.clone()
}
#[cfg(test)]
pub(crate) fn stream_waker_entries(&self) -> (usize, usize) {
self.cell.borrow().stream_waker_entries()
}
#[cfg(test)]
pub(crate) fn sugar_waker_entries(&self) -> (usize, usize, usize) {
self.cell.borrow().sugar_waker_entries()
}
pub fn is_established(&self) -> bool {
self.cell.borrow().is_established()
}
}
impl<S: Handshake + 'static> Connection<S> {
#[allow(dead_code)]
pub(crate) fn notification_slot_boxed(&self) -> WakerSlot<Box<dyn FnMut(u64)>> {
let key = self.cell.borrow_mut().notification_wakers.key();
WakerSlot::new(
key,
Box::new({
let cell = Rc::clone(&self.cell);
move |key| cell.borrow_mut().notification_wakers.unpark(key)
}),
)
}
#[allow(dead_code)]
pub(crate) fn settled_slot_boxed(&self) -> WakerSlot<Box<dyn FnMut(u64)>> {
let key = self.cell.borrow_mut().settled_wakers.key();
WakerSlot::new(
key,
Box::new({
let cell = Rc::clone(&self.cell);
move |key| cell.borrow_mut().settled_wakers.unpark(key)
}),
)
}
#[allow(dead_code)]
pub(crate) fn message_reader_slot_boxed(&self) -> WakerSlot<Box<dyn FnMut(u64)>> {
let key = self.cell.borrow_mut().message_readers.key();
WakerSlot::new(
key,
Box::new({
let cell = Rc::clone(&self.cell);
move |key| cell.borrow_mut().message_readers.unpark(key)
}),
)
}
#[allow(dead_code)]
pub(crate) fn datagram_reader_slot_boxed(&self) -> WakerSlot<Box<dyn FnMut(u64)>> {
let key = self.cell.borrow_mut().datagram_readers.key();
WakerSlot::new(
key,
Box::new({
let cell = Rc::clone(&self.cell);
move |key| cell.borrow_mut().datagram_readers.unpark(key)
}),
)
}
#[allow(dead_code)]
pub(crate) fn message_sender_slot_boxed(&self) -> WakerSlot<Box<dyn FnMut(u64)>> {
let key = self.cell.borrow_mut().message_senders.key();
WakerSlot::new(
key,
Box::new({
let cell = Rc::clone(&self.cell);
move |key| cell.borrow_mut().message_senders.unpark(key)
}),
)
}
#[allow(dead_code)]
pub(crate) fn opener_slot_boxed(&self, dir: Dir) -> WakerSlot<Box<dyn FnMut(u64)>> {
let key = self.cell.borrow_mut().stream_openers[dir.slot()].key();
WakerSlot::new(
key,
Box::new({
let cell = Rc::clone(&self.cell);
move |key| cell.borrow_mut().stream_openers[dir.slot()].unpark(key)
}),
)
}
#[allow(dead_code)]
pub(crate) fn acceptor_slot_boxed(&self, dir: Dir) -> WakerSlot<Box<dyn FnMut(u64)>> {
let key = self.cell.borrow_mut().stream_acceptors[dir.slot()].key();
WakerSlot::new(
key,
Box::new({
let cell = Rc::clone(&self.cell);
move |key| cell.borrow_mut().stream_acceptors[dir.slot()].unpark(key)
}),
)
}
}
fn install_bi<S: Handshake>(
shell: Rc<dyn ShellLink>,
cell_rc: Rc<RefCell<ConnCell<S>>>,
cell: &mut ConnCell<S>,
conn: ConnectionId,
r: StreamRef,
) -> BiStream<S> {
let send = SendStream::install(Rc::clone(&shell), Rc::clone(&cell_rc), cell, conn, r);
let recv = RecvStream::install(shell, cell_rc, cell, conn, r);
BiStream::new(send, recv)
}
impl<S: Handshake> std::fmt::Debug for Connection<S> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Connection")
.field("id", &self.id)
.field("remote_address", &self.remote_address())
.field("established", &self.is_established())
.finish_non_exhaustive()
}
}
impl<S: Handshake> Drop for Connection<S> {
fn drop(&mut self) {
let last_for_connection = {
let mut cell = self.cell.borrow_mut();
cell.handles -= 1;
cell.handles == 0
};
let last_in_process = self.shell.release();
if last_for_connection && !last_in_process {
self.close_now(constants::NO_ERROR, b"");
}
}
}