use std::{
mem::replace,
net::SocketAddr,
pin::Pin,
sync::{Arc, RwLock},
task::{Context, Poll},
};
use async_compat::Compat;
use dht_rpc::Rpc;
use futures::{Sink, Stream};
use hypercore_handshake::{Cipher, CipherEvent, CipherIo, CipherTrait};
use udx::HalfOpenStreamHandle;
use uint24le_framing::Uint24LELengthPrefixedFraming;
use crate::Error;
#[derive(Debug)]
pub enum ConnStep {
Start(HalfOpenStreamHandle),
Ready,
Failed,
}
pub struct ConnectionInner {
pub handshake: Cipher,
pub udx_local_id: u32,
pub step: ConnStep,
pub rpc: Option<Rpc>,
}
impl std::fmt::Debug for ConnectionInner {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ConnectionInner")
.field("handshake", &self.handshake)
.field("udx_local_id", &self.udx_local_id)
.field("step", &self.step)
.finish()
}
}
impl ConnectionInner {
pub fn new(handshake: Cipher, udx_local_id: u32, half_stream: HalfOpenStreamHandle) -> Self {
Self {
handshake,
udx_local_id,
step: ConnStep::Start(half_stream),
rpc: None,
}
}
pub fn new_with_rpc(
handshake: Cipher,
udx_local_id: u32,
half_stream: HalfOpenStreamHandle,
rpc: Rpc,
) -> Self {
Self {
handshake,
udx_local_id,
step: ConnStep::Start(half_stream),
rpc: Some(rpc),
}
}
pub fn receive_next(&mut self, noise: Vec<u8>) -> Result<CipherEvent, Error> {
self.handshake.receive_next(noise);
Ok(self
.handshake
.next_decrypted_message()?
.expect("recieved msg above"))
}
pub fn handshake_ready(&self) -> bool {
self.handshake.ready()
}
pub fn udx_local_id(&self) -> u32 {
self.udx_local_id
}
pub fn handshake_set_io(&mut self, io: Box<dyn CipherIo<Error = std::io::Error>>) {
self.handshake.set_io(io)
}
pub fn set_step(&mut self, step: ConnStep) {
self.step = step;
}
fn connect(&mut self, addr: SocketAddr, remote_id: u32) -> Result<(), Error> {
let ConnStep::Start(half_stream) = replace(&mut self.step, ConnStep::Failed) else {
todo!()
};
let stream = half_stream.connect(addr, remote_id)?;
let framed_udx_stream = Uint24LELengthPrefixedFraming::new(Compat::new(stream.clone()));
self.handshake_set_io(Box::new(framed_udx_stream));
self.step = ConnStep::Ready;
Ok(())
}
}
impl Stream for ConnectionInner {
type Item = CipherEvent;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
if let Some(ref mut rpc) = self.rpc {
while Pin::new(&mut *rpc).poll_next(cx).is_ready() {
}
}
Pin::new(&mut self.handshake).poll_next(cx)
}
}
impl Sink<Vec<u8>> for ConnectionInner {
type Error = std::io::Error;
fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.handshake).poll_ready(cx)
}
fn start_send(mut self: Pin<&mut Self>, item: Vec<u8>) -> Result<(), Self::Error> {
Pin::new(&mut self.handshake).start_send(item)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.handshake).poll_flush(cx)
}
fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.handshake).poll_close(cx)
}
}
#[derive(Clone)]
pub struct Connection {
pub inner: Arc<RwLock<ConnectionInner>>,
}
impl std::fmt::Debug for Connection {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self.inner.try_read() {
Ok(x) => f.debug_struct("Connection").field("inner", &x).finish(),
Err(_) => f.debug_struct("Connection").field("inner", &()).finish(),
}
}
}
macro_rules! w {
($self:expr) => {
$self.inner.write().unwrap()
};
}
macro_rules! r {
($self:expr) => {
$self.inner.read().unwrap()
};
}
impl Connection {
pub fn new(handshake: Cipher, udx_local_id: u32, half_stream: HalfOpenStreamHandle) -> Self {
Self {
inner: Arc::new(RwLock::new(ConnectionInner {
handshake,
udx_local_id,
step: ConnStep::Start(half_stream),
rpc: None,
})),
}
}
pub fn new_with_rpc(
handshake: Cipher,
udx_local_id: u32,
half_stream: HalfOpenStreamHandle,
rpc: Rpc,
) -> Self {
Self {
inner: Arc::new(RwLock::new(ConnectionInner {
handshake,
udx_local_id,
step: ConnStep::Start(half_stream),
rpc: Some(rpc),
})),
}
}
pub fn receive_next(&self, noise: Vec<u8>) -> Result<CipherEvent, Error> {
w!(self).handshake.receive_next(noise);
Ok(w!(self)
.handshake
.next_decrypted_message()?
.expect("recieved msg above"))
}
pub fn step_ready(&self) -> bool {
matches!(r!(self).step, ConnStep::Ready)
}
pub fn handshake_ready(&self) -> bool {
r!(self).handshake.ready()
}
pub fn udx_local_id(&self) -> u32 {
r!(self).udx_local_id
}
pub fn handshake_set_io(&self, io: Box<dyn CipherIo<Error = std::io::Error>>) {
w!(self).handshake.set_io(io)
}
pub fn set_step(&self, step: ConnStep) {
w!(self).step = step;
}
pub fn connect(&self, addr: SocketAddr, remote_id: u32) -> Result<(), Error> {
w!(self).connect(addr, remote_id)
}
pub fn get_remote_static(&self) -> Option<[u8; 32]> {
r!(self).handshake.get_remote_static()
}
pub fn handshake_hash(&self) -> Option<Vec<u8>> {
r!(self).handshake.handshake_hash().map(|h| h.to_vec())
}
}
impl Stream for Connection {
type Item = CipherEvent;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let mut inner = self.inner.write().unwrap();
if let Some(ref mut rpc) = inner.rpc {
while Pin::new(&mut *rpc).poll_next(cx).is_ready() {
}
}
Pin::new(&mut inner.handshake).poll_next(cx)
}
}
impl Sink<Vec<u8>> for Connection {
type Error = std::io::Error;
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.inner.write().unwrap().handshake).poll_ready(cx)
}
fn start_send(self: Pin<&mut Self>, item: Vec<u8>) -> Result<(), Self::Error> {
Pin::new(&mut self.inner.write().unwrap().handshake).start_send(item)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.inner.write().unwrap().handshake).poll_flush(cx)
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.inner.write().unwrap().handshake).poll_close(cx)
}
}
impl CipherTrait for Connection {
fn remote_public_key(&self) -> Option<[u8; hypercore_handshake::state_machine::PUBLIC_KEYLEN]> {
self.inner.read().unwrap().handshake.get_remote_static()
}
fn local_public_key(&self) -> [u8; hypercore_handshake::state_machine::PUBLIC_KEYLEN] {
self.inner.read().unwrap().handshake.get_local_public_key()
}
fn handshake_hash(&self) -> Option<Vec<u8>> {
self.inner
.read()
.unwrap()
.handshake
.handshake_hash()
.map(|h| h.to_vec())
}
fn is_initiator(&self) -> bool {
self.inner.read().unwrap().handshake.is_initiator()
}
}