2020-12-02 12:30:47 +00:00
|
|
|
use std::collections::HashMap;
|
|
|
|
use std::net::SocketAddr;
|
2021-10-12 15:59:46 +00:00
|
|
|
use std::sync::atomic::{self, AtomicBool, AtomicU32};
|
2020-12-07 12:35:24 +00:00
|
|
|
use std::sync::{Arc, Mutex};
|
2020-12-02 12:30:47 +00:00
|
|
|
|
|
|
|
use bytes::Bytes;
|
2020-12-07 12:35:24 +00:00
|
|
|
use log::{debug, error, trace};
|
2020-12-02 12:30:47 +00:00
|
|
|
|
|
|
|
use tokio::net::TcpStream;
|
|
|
|
use tokio::sync::{mpsc, oneshot, watch};
|
2021-10-12 12:51:28 +00:00
|
|
|
use tokio_util::compat::*;
|
|
|
|
|
|
|
|
use futures::io::AsyncReadExt;
|
2020-12-02 12:30:47 +00:00
|
|
|
|
2020-12-07 12:35:24 +00:00
|
|
|
use async_trait::async_trait;
|
|
|
|
|
2021-10-12 12:56:29 +00:00
|
|
|
use kuska_handshake::async_std::{handshake_client, handshake_server, BoxStream};
|
2020-12-02 12:30:47 +00:00
|
|
|
|
2021-10-12 15:59:46 +00:00
|
|
|
use crate::endpoint::*;
|
2020-12-02 12:30:47 +00:00
|
|
|
use crate::error::*;
|
|
|
|
use crate::netapp::*;
|
|
|
|
use crate::proto::*;
|
|
|
|
use crate::util::*;
|
|
|
|
|
2021-10-12 15:59:46 +00:00
|
|
|
// Request message format (client -> server):
|
|
|
|
// - u8 priority
|
|
|
|
// - u8 path length
|
|
|
|
// - [u8; path length] path
|
|
|
|
// - [u8; *] data
|
|
|
|
|
|
|
|
// Response message format (server -> client):
|
|
|
|
// - u8 response code
|
|
|
|
// - [u8; *] response
|
|
|
|
|
2020-12-02 19:12:24 +00:00
|
|
|
pub(crate) struct ServerConn {
|
|
|
|
pub(crate) remote_addr: SocketAddr,
|
2020-12-12 20:14:15 +00:00
|
|
|
pub(crate) peer_id: NodeID,
|
2020-12-02 19:12:24 +00:00
|
|
|
|
2020-12-02 12:30:47 +00:00
|
|
|
netapp: Arc<NetApp>,
|
2020-12-02 19:12:24 +00:00
|
|
|
|
2020-12-07 12:35:24 +00:00
|
|
|
resp_send: mpsc::UnboundedSender<Option<(RequestID, RequestPriority, Vec<u8>)>>,
|
2020-12-02 12:30:47 +00:00
|
|
|
close_send: watch::Sender<bool>,
|
|
|
|
}
|
|
|
|
|
|
|
|
impl ServerConn {
|
|
|
|
pub(crate) async fn run(netapp: Arc<NetApp>, socket: TcpStream) -> Result<(), Error> {
|
2021-10-12 12:51:28 +00:00
|
|
|
let remote_addr = socket.peer_addr()?;
|
|
|
|
let mut socket = socket.compat();
|
|
|
|
|
2020-12-02 12:30:47 +00:00
|
|
|
let handshake = handshake_server(
|
2021-10-12 12:51:28 +00:00
|
|
|
&mut socket,
|
2020-12-02 12:30:47 +00:00
|
|
|
netapp.netid.clone(),
|
2021-10-12 11:18:24 +00:00
|
|
|
netapp.id,
|
2020-12-02 12:30:47 +00:00
|
|
|
netapp.privkey.clone(),
|
|
|
|
)
|
|
|
|
.await?;
|
2021-10-12 11:18:24 +00:00
|
|
|
let peer_id = handshake.peer_pk;
|
2020-12-02 12:30:47 +00:00
|
|
|
|
|
|
|
debug!(
|
|
|
|
"Handshake complete (server) with {}@{}",
|
2020-12-12 20:14:15 +00:00
|
|
|
hex::encode(&peer_id),
|
2020-12-02 12:30:47 +00:00
|
|
|
remote_addr
|
|
|
|
);
|
|
|
|
|
2021-10-12 12:51:28 +00:00
|
|
|
let (read, write) = socket.split();
|
2020-12-02 12:30:47 +00:00
|
|
|
|
2020-12-07 17:07:55 +00:00
|
|
|
let (read, write) =
|
2020-12-02 12:30:47 +00:00
|
|
|
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,
|
2021-10-12 11:18:24 +00:00
|
|
|
peer_id,
|
2020-12-02 12:30:47 +00:00
|
|
|
resp_send,
|
|
|
|
close_send,
|
|
|
|
});
|
|
|
|
|
2021-10-12 11:18:24 +00:00
|
|
|
netapp.connected_as_server(peer_id, conn.clone());
|
2020-12-02 12:30:47 +00:00
|
|
|
|
|
|
|
let conn2 = conn.clone();
|
|
|
|
let conn3 = conn.clone();
|
2020-12-07 12:35:24 +00:00
|
|
|
let close_recv2 = close_recv.clone();
|
2020-12-02 12:30:47 +00:00
|
|
|
tokio::try_join!(
|
2020-12-07 12:35:24 +00:00
|
|
|
async move {
|
|
|
|
tokio::select!(
|
2020-12-07 17:07:55 +00:00
|
|
|
r = conn2.recv_loop(read) => r,
|
2020-12-07 12:35:24 +00:00
|
|
|
_ = await_exit(close_recv) => Ok(()),
|
|
|
|
)
|
|
|
|
},
|
|
|
|
async move {
|
|
|
|
tokio::select!(
|
2020-12-07 17:07:55 +00:00
|
|
|
r = conn3.send_loop(resp_recv, write) => r,
|
2020-12-07 12:35:24 +00:00
|
|
|
_ = await_exit(close_recv2) => Ok(()),
|
|
|
|
)
|
|
|
|
},
|
2020-12-02 12:30:47 +00:00
|
|
|
)
|
|
|
|
.map(|_| ())
|
|
|
|
.log_err("ServerConn recv_loop/send_loop");
|
|
|
|
|
2020-12-12 20:14:15 +00:00
|
|
|
netapp.disconnected_as_server(&peer_id, conn);
|
2020-12-02 12:30:47 +00:00
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
|
|
|
pub fn close(&self) {
|
2021-10-12 11:07:34 +00:00
|
|
|
self.close_send.send(true).unwrap();
|
2020-12-02 12:30:47 +00:00
|
|
|
}
|
2021-10-12 15:59:46 +00:00
|
|
|
|
|
|
|
async fn recv_handler_aux(self: &Arc<Self>, bytes: &[u8]) -> Result<Vec<u8>, Error> {
|
|
|
|
if bytes.len() < 2 {
|
|
|
|
return Err(Error::Message("Invalid protocol message".into()));
|
|
|
|
}
|
|
|
|
|
|
|
|
// byte 0 is the request priority, we don't care here
|
|
|
|
let path_length = bytes[1] as usize;
|
|
|
|
if bytes.len() < 2 + path_length {
|
|
|
|
return Err(Error::Message("Invalid protocol message".into()));
|
|
|
|
}
|
|
|
|
|
|
|
|
let path = &bytes[2..2 + path_length];
|
|
|
|
let path = String::from_utf8(path.to_vec())?;
|
|
|
|
let data = &bytes[2 + path_length..];
|
|
|
|
|
|
|
|
let handler_opt = {
|
|
|
|
let endpoints = self.netapp.endpoints.read().unwrap();
|
|
|
|
endpoints.get(&path).map(|e| e.clone_endpoint())
|
|
|
|
};
|
|
|
|
|
|
|
|
if let Some(handler) = handler_opt {
|
|
|
|
handler.handle(data, self.peer_id).await
|
|
|
|
} else {
|
|
|
|
Err(Error::NoHandler)
|
|
|
|
}
|
|
|
|
}
|
2020-12-02 12:30:47 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
impl SendLoop for ServerConn {}
|
|
|
|
|
|
|
|
#[async_trait]
|
|
|
|
impl RecvLoop for ServerConn {
|
2021-10-12 15:59:46 +00:00
|
|
|
async fn recv_handler(self: Arc<Self>, id: RequestID, bytes: Vec<u8>) {
|
2020-12-11 12:34:04 +00:00
|
|
|
trace!("ServerConn recv_handler {} ({} bytes)", id, bytes.len());
|
2020-12-02 12:30:47 +00:00
|
|
|
let bytes: Bytes = bytes.into();
|
|
|
|
|
2021-10-12 15:59:46 +00:00
|
|
|
let resp = self.recv_handler_aux(&bytes[..]).await;
|
2020-12-02 12:30:47 +00:00
|
|
|
let prio = bytes[0];
|
|
|
|
|
2021-10-12 15:59:46 +00:00
|
|
|
let mut resp_bytes = vec![];
|
|
|
|
match resp {
|
|
|
|
Ok(rb) => {
|
|
|
|
resp_bytes.push(0u8);
|
|
|
|
resp_bytes.extend(&rb[..]);
|
|
|
|
}
|
|
|
|
Err(e) => {
|
|
|
|
resp_bytes.push(e.code());
|
|
|
|
}
|
2020-12-02 12:30:47 +00:00
|
|
|
}
|
2021-10-12 15:59:46 +00:00
|
|
|
|
|
|
|
self.resp_send
|
|
|
|
.send(Some((id, prio, resp_bytes)))
|
|
|
|
.log_err("ServerConn recv_handler send resp");
|
2020-12-02 12:30:47 +00:00
|
|
|
}
|
|
|
|
}
|
2020-12-02 19:12:24 +00:00
|
|
|
pub(crate) struct ClientConn {
|
|
|
|
pub(crate) remote_addr: SocketAddr,
|
2020-12-12 20:14:15 +00:00
|
|
|
pub(crate) peer_id: NodeID,
|
2020-12-02 19:12:24 +00:00
|
|
|
|
2020-12-07 12:35:24 +00:00
|
|
|
query_send: mpsc::UnboundedSender<Option<(RequestID, RequestPriority, Vec<u8>)>>,
|
|
|
|
|
2021-10-12 15:59:46 +00:00
|
|
|
next_query_number: AtomicU32,
|
2020-12-07 12:35:24 +00:00
|
|
|
inflight: Mutex<HashMap<RequestID, oneshot::Sender<Vec<u8>>>>,
|
|
|
|
must_exit: AtomicBool,
|
|
|
|
stop_recv_loop: watch::Sender<bool>,
|
2020-12-02 12:30:47 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
impl ClientConn {
|
|
|
|
pub(crate) async fn init(
|
|
|
|
netapp: Arc<NetApp>,
|
|
|
|
socket: TcpStream,
|
2020-12-12 20:14:15 +00:00
|
|
|
peer_id: NodeID,
|
2020-12-02 12:30:47 +00:00
|
|
|
) -> Result<(), Error> {
|
2021-10-12 12:51:28 +00:00
|
|
|
let remote_addr = socket.peer_addr()?;
|
|
|
|
let mut socket = socket.compat();
|
2020-12-02 12:30:47 +00:00
|
|
|
|
|
|
|
let handshake = handshake_client(
|
2021-10-12 12:51:28 +00:00
|
|
|
&mut socket,
|
2020-12-02 12:30:47 +00:00
|
|
|
netapp.netid.clone(),
|
2021-10-12 11:18:24 +00:00
|
|
|
netapp.id,
|
2020-12-02 12:30:47 +00:00
|
|
|
netapp.privkey.clone(),
|
2021-10-12 11:18:24 +00:00
|
|
|
peer_id,
|
2020-12-02 12:30:47 +00:00
|
|
|
)
|
|
|
|
.await?;
|
|
|
|
|
|
|
|
debug!(
|
|
|
|
"Handshake complete (client) with {}@{}",
|
2020-12-12 20:14:15 +00:00
|
|
|
hex::encode(&peer_id),
|
2020-12-02 12:30:47 +00:00
|
|
|
remote_addr
|
|
|
|
);
|
|
|
|
|
2021-10-12 12:51:28 +00:00
|
|
|
let (read, write) = socket.split();
|
2020-12-02 12:30:47 +00:00
|
|
|
|
2020-12-07 17:07:55 +00:00
|
|
|
let (read, write) =
|
2020-12-02 12:30:47 +00:00
|
|
|
BoxStream::from_handshake(read, write, handshake, 0x8000).split_read_write();
|
|
|
|
|
|
|
|
let (query_send, query_recv) = mpsc::unbounded_channel();
|
|
|
|
|
2020-12-07 12:35:24 +00:00
|
|
|
let (stop_recv_loop, stop_recv_loop_recv) = watch::channel(false);
|
2020-12-02 12:30:47 +00:00
|
|
|
|
|
|
|
let conn = Arc::new(ClientConn {
|
|
|
|
remote_addr,
|
2021-10-12 11:18:24 +00:00
|
|
|
peer_id,
|
2021-10-12 15:59:46 +00:00
|
|
|
next_query_number: AtomicU32::from(RequestID::default()),
|
2020-12-02 12:30:47 +00:00
|
|
|
query_send,
|
2020-12-07 12:35:24 +00:00
|
|
|
inflight: Mutex::new(HashMap::new()),
|
|
|
|
must_exit: AtomicBool::new(false),
|
|
|
|
stop_recv_loop,
|
2020-12-02 12:30:47 +00:00
|
|
|
});
|
|
|
|
|
2021-10-12 11:18:24 +00:00
|
|
|
netapp.connected_as_client(peer_id, conn.clone());
|
2020-12-02 12:30:47 +00:00
|
|
|
|
|
|
|
tokio::spawn(async move {
|
|
|
|
let conn2 = conn.clone();
|
|
|
|
let conn3 = conn.clone();
|
2020-12-07 17:07:55 +00:00
|
|
|
tokio::try_join!(conn2.send_loop(query_recv, write), async move {
|
|
|
|
tokio::select!(
|
|
|
|
r = conn3.recv_loop(read) => r,
|
|
|
|
_ = await_exit(stop_recv_loop_recv) => Ok(()),
|
|
|
|
)
|
|
|
|
})
|
2020-12-02 12:30:47 +00:00
|
|
|
.map(|_| ())
|
|
|
|
.log_err("ClientConn send_loop/recv_loop/dispatch_loop");
|
|
|
|
|
2020-12-12 20:14:15 +00:00
|
|
|
netapp.disconnected_as_client(&peer_id, conn);
|
2020-12-02 12:30:47 +00:00
|
|
|
});
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
|
|
|
pub fn close(&self) {
|
2020-12-07 12:35:24 +00:00
|
|
|
self.must_exit.store(true, atomic::Ordering::SeqCst);
|
|
|
|
self.query_send
|
|
|
|
.send(None)
|
|
|
|
.log_err("could not write None in query_send");
|
|
|
|
if self.inflight.lock().unwrap().is_empty() {
|
|
|
|
self.stop_recv_loop
|
2021-10-12 11:07:34 +00:00
|
|
|
.send(true)
|
2020-12-07 12:35:24 +00:00
|
|
|
.log_err("could not write true to stop_recv_loop");
|
2020-12-02 12:30:47 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2021-10-12 15:59:46 +00:00
|
|
|
pub(crate) async fn call<T>(
|
2020-12-02 12:30:47 +00:00
|
|
|
self: Arc<Self>,
|
|
|
|
rq: T,
|
2021-10-12 15:59:46 +00:00
|
|
|
path: &str,
|
2020-12-02 12:30:47 +00:00
|
|
|
prio: RequestPriority,
|
|
|
|
) -> Result<<T as Message>::Response, Error>
|
|
|
|
where
|
|
|
|
T: Message,
|
|
|
|
{
|
|
|
|
let id = self
|
|
|
|
.next_query_number
|
2021-10-12 15:59:46 +00:00
|
|
|
.fetch_add(1, atomic::Ordering::Relaxed);
|
|
|
|
let mut bytes = vec![prio, path.as_bytes().len() as u8];
|
|
|
|
bytes.extend_from_slice(path.as_bytes());
|
2020-12-02 12:30:47 +00:00
|
|
|
bytes.extend_from_slice(&rmp_to_vec_all_named(&rq)?[..]);
|
|
|
|
|
|
|
|
let (resp_send, resp_recv) = oneshot::channel();
|
2020-12-07 12:35:24 +00:00
|
|
|
let old = self.inflight.lock().unwrap().insert(id, resp_send);
|
|
|
|
if let Some(old_ch) = old {
|
|
|
|
error!(
|
|
|
|
"Too many inflight requests! RequestID collision. Interrupting previous request."
|
|
|
|
);
|
|
|
|
if old_ch.send(vec![]).is_err() {
|
|
|
|
debug!("Could not send empty response to collisionned request, probably because request was interrupted. Dropping response.");
|
|
|
|
}
|
|
|
|
}
|
2020-12-02 12:30:47 +00:00
|
|
|
|
|
|
|
trace!("request: query_send {}, {} bytes", id, bytes.len());
|
2020-12-07 12:35:24 +00:00
|
|
|
self.query_send.send(Some((id, prio, bytes)))?;
|
2020-12-02 12:30:47 +00:00
|
|
|
|
|
|
|
let resp = resp_recv.await?;
|
|
|
|
|
2021-10-12 15:59:46 +00:00
|
|
|
let code = resp[0];
|
|
|
|
if code == 0 {
|
|
|
|
rmp_serde::decode::from_read_ref::<_, Result<<T as Message>::Response, String>>(
|
|
|
|
&resp[1..],
|
|
|
|
)?
|
2020-12-02 12:30:47 +00:00
|
|
|
.map_err(Error::Remote)
|
2021-10-12 15:59:46 +00:00
|
|
|
} else {
|
|
|
|
Err(Error::Remote(format!("Remote error code {}", code)))
|
|
|
|
}
|
2020-12-02 12:30:47 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl SendLoop for ClientConn {}
|
|
|
|
|
|
|
|
#[async_trait]
|
|
|
|
impl RecvLoop for ClientConn {
|
|
|
|
async fn recv_handler(self: Arc<Self>, id: RequestID, msg: Vec<u8>) {
|
2020-12-11 12:34:04 +00:00
|
|
|
trace!("ClientConn recv_handler {} ({} bytes)", id, msg.len());
|
2020-12-07 15:00:12 +00:00
|
|
|
|
2020-12-07 12:35:24 +00:00
|
|
|
let mut inflight = self.inflight.lock().unwrap();
|
|
|
|
if let Some(ch) = inflight.remove(&id) {
|
|
|
|
if ch.send(msg).is_err() {
|
|
|
|
debug!("Could not send request response, probably because request was interrupted. Dropping response.");
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
if inflight.is_empty() && self.must_exit.load(atomic::Ordering::SeqCst) {
|
|
|
|
self.stop_recv_loop
|
2021-10-12 11:07:34 +00:00
|
|
|
.send(true)
|
2020-12-07 12:35:24 +00:00
|
|
|
.log_err("could not write true to stop_recv_loop");
|
|
|
|
}
|
2020-12-02 12:30:47 +00:00
|
|
|
}
|
|
|
|
}
|