netapp/src/conn.rs

282 lines
7.4 KiB
Rust
Raw Normal View History

2020-12-02 12:30:47 +00:00
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::atomic::{self, AtomicU16};
use std::sync::Arc;
use async_trait::async_trait;
use bytes::Bytes;
use log::{debug, trace};
use sodiumoxide::crypto::sign::ed25519;
use tokio::io::split;
use tokio::net::TcpStream;
use tokio::sync::{mpsc, oneshot, watch};
use kuska_handshake::async_std::{
handshake_client, handshake_server, BoxStream, TokioCompatExt, TokioCompatExtRead,
TokioCompatExtWrite,
};
use crate::error::*;
use crate::message::*;
use crate::netapp::*;
use crate::proto::*;
use crate::util::*;
pub struct ServerConn {
netapp: Arc<NetApp>,
pub remote_addr: SocketAddr,
pub peer_pk: ed25519::PublicKey,
resp_send: mpsc::UnboundedSender<(RequestID, RequestPriority, Vec<u8>)>,
close_send: watch::Sender<bool>,
}
impl ServerConn {
pub(crate) async fn run(netapp: Arc<NetApp>, socket: TcpStream) -> Result<(), Error> {
let mut asyncstd_socket = TokioCompatExt::wrap(socket);
let handshake = handshake_server(
&mut asyncstd_socket,
netapp.netid.clone(),
netapp.pubkey.clone(),
netapp.privkey.clone(),
)
.await?;
let peer_pk = handshake.peer_pk.clone();
let tokio_socket = asyncstd_socket.into_inner();
let remote_addr = tokio_socket.peer_addr().unwrap();
debug!(
"Handshake complete (server) with {}@{}",
hex::encode(&peer_pk),
remote_addr
);
let (read, write) = split(tokio_socket);
let read = TokioCompatExtRead::wrap(read);
let write = TokioCompatExtWrite::wrap(write);
let (box_stream_read, box_stream_write) =
BoxStream::from_handshake(read, write, handshake, 0x8000).split_read_write();
let (resp_send, resp_recv) = mpsc::unbounded_channel();
let (close_send, close_recv) = watch::channel(false);
let conn = Arc::new(ServerConn {
netapp: netapp.clone(),
remote_addr,
peer_pk: peer_pk.clone(),
resp_send,
close_send,
});
netapp.connected_as_server(peer_pk.clone(), conn.clone());
let conn2 = conn.clone();
let conn3 = conn.clone();
tokio::try_join!(
conn2.recv_loop(box_stream_read, close_recv.clone()),
conn3.send_loop(resp_recv, box_stream_write, close_recv.clone()),
)
.map(|_| ())
.log_err("ServerConn recv_loop/send_loop");
netapp.disconnected_as_server(&peer_pk, conn);
Ok(())
}
pub fn close(&self) {
self.close_send.broadcast(true).unwrap();
}
}
impl SendLoop for ServerConn {}
#[async_trait]
impl RecvLoop for ServerConn {
async fn recv_handler(self: Arc<Self>, id: u16, bytes: Vec<u8>) {
let bytes: Bytes = bytes.into();
let prio = bytes[0];
let mut kind_bytes = [0u8; 4];
kind_bytes.copy_from_slice(&bytes[1..5]);
let kind = u32::from_be_bytes(kind_bytes);
if let Some(handler) = self.netapp.msg_handlers.load().get(&kind) {
2020-12-02 17:10:07 +00:00
let net_handler = &handler.net_handler;
let resp = net_handler(self.peer_pk.clone(), bytes.slice(5..)).await;
2020-12-02 12:30:47 +00:00
self.resp_send
.send((id, prio, resp))
.log_err("ServerConn recv_handler send resp");
}
}
}
pub struct ClientConn {
pub netapp: Arc<NetApp>,
pub remote_addr: SocketAddr,
pub peer_pk: ed25519::PublicKey,
query_send: mpsc::UnboundedSender<(RequestID, RequestPriority, Vec<u8>)>,
next_query_number: AtomicU16,
resp_send: mpsc::UnboundedSender<(RequestID, Vec<u8>)>,
resp_notify_send: mpsc::UnboundedSender<(RequestID, oneshot::Sender<Vec<u8>>)>,
close_send: watch::Sender<bool>,
}
impl ClientConn {
pub(crate) async fn init(
netapp: Arc<NetApp>,
socket: TcpStream,
remote_pk: ed25519::PublicKey,
) -> Result<(), Error> {
let mut asyncstd_socket = TokioCompatExt::wrap(socket);
let handshake = handshake_client(
&mut asyncstd_socket,
netapp.netid.clone(),
netapp.pubkey.clone(),
netapp.privkey.clone(),
remote_pk.clone(),
)
.await?;
let tokio_socket = asyncstd_socket.into_inner();
let remote_addr = tokio_socket.peer_addr().unwrap();
debug!(
"Handshake complete (client) with {}@{}",
hex::encode(&remote_pk),
remote_addr
);
let (read, write) = split(tokio_socket);
let read = TokioCompatExtRead::wrap(read);
let write = TokioCompatExtWrite::wrap(write);
let (box_stream_read, box_stream_write) =
BoxStream::from_handshake(read, write, handshake, 0x8000).split_read_write();
let (query_send, query_recv) = mpsc::unbounded_channel();
let (resp_send, resp_recv) = mpsc::unbounded_channel();
let (resp_notify_send, resp_notify_recv) = mpsc::unbounded_channel();
let (close_send, close_recv) = watch::channel(false);
let conn = Arc::new(ClientConn {
netapp: netapp.clone(),
remote_addr,
peer_pk: remote_pk.clone(),
next_query_number: AtomicU16::from(0u16),
query_send,
resp_send,
resp_notify_send,
close_send,
});
netapp.connected_as_client(remote_pk.clone(), conn.clone());
tokio::spawn(async move {
let conn2 = conn.clone();
let conn3 = conn.clone();
let conn4 = conn.clone();
tokio::try_join!(
conn2.send_loop(query_recv, box_stream_write, close_recv.clone()),
conn3.recv_loop(box_stream_read, close_recv.clone()),
conn4.dispatch_resp(resp_recv, resp_notify_recv, close_recv.clone()),
)
.map(|_| ())
.log_err("ClientConn send_loop/recv_loop/dispatch_loop");
netapp.disconnected_as_client(&remote_pk, conn);
});
Ok(())
}
pub fn close(&self) {
self.close_send.broadcast(true).unwrap();
}
async fn dispatch_resp(
self: Arc<Self>,
mut resp_recv: mpsc::UnboundedReceiver<(RequestID, Vec<u8>)>,
mut resp_notify_recv: mpsc::UnboundedReceiver<(RequestID, oneshot::Sender<Vec<u8>>)>,
mut must_exit: watch::Receiver<bool>,
) -> Result<(), Error> {
let mut resps: HashMap<RequestID, Vec<u8>> = HashMap::new();
let mut resp_notify: HashMap<RequestID, oneshot::Sender<Vec<u8>>> = HashMap::new();
while !*must_exit.borrow() {
tokio::select! {
resp = resp_recv.recv() => {
if let Some((id, resp)) = resp {
trace!("dispatch_resp: got resp to {}, {} bytes", id, resp.len());
if let Some(ch) = resp_notify.remove(&id) {
ch.send(resp).map_err(|_| Error::Message("Could not dispatch reply".to_string()))?;
} else {
resps.insert(id, resp);
}
}
}
resp_ch = resp_notify_recv.recv() => {
if let Some((id, resp_ch)) = resp_ch {
trace!("dispatch_resp: got resp_ch {}", id);
if let Some(rs) = resps.remove(&id) {
resp_ch.send(rs).map_err(|_| Error::Message("Could not dispatch reply".to_string()))?;
} else {
resp_notify.insert(id, resp_ch);
}
}
}
exit = must_exit.recv() => {
if exit == Some(true) {
break;
}
}
}
}
Ok(())
}
2020-12-02 17:10:07 +00:00
pub(crate) async fn request<T>(
2020-12-02 12:30:47 +00:00
self: Arc<Self>,
rq: T,
prio: RequestPriority,
) -> Result<<T as Message>::Response, Error>
where
T: Message,
{
let id = self
.next_query_number
.fetch_add(1u16, atomic::Ordering::Relaxed);
let mut bytes = vec![prio];
bytes.extend_from_slice(&u32::to_be_bytes(T::KIND)[..]);
bytes.extend_from_slice(&rmp_to_vec_all_named(&rq)?[..]);
let (resp_send, resp_recv) = oneshot::channel();
self.resp_notify_send.send((id, resp_send))?;
trace!("request: query_send {}, {} bytes", id, bytes.len());
self.query_send.send((id, prio, bytes))?;
let resp = resp_recv.await?;
rmp_serde::decode::from_read_ref::<_, Result<<T as Message>::Response, String>>(&resp[..])?
.map_err(Error::Remote)
}
}
impl SendLoop for ClientConn {}
#[async_trait]
impl RecvLoop for ClientConn {
async fn recv_handler(self: Arc<Self>, id: RequestID, msg: Vec<u8>) {
self.resp_send
.send((id, msg))
.log_err("ClientConn::recv_handler");
}
}