netapp/src/server.rs

207 lines
5 KiB
Rust
Raw Normal View History

use std::net::SocketAddr;
use std::sync::Arc;
use arc_swap::ArcSwapOption;
use bytes::Bytes;
use log::{debug, trace};
2022-02-18 18:01:23 +00:00
#[cfg(feature = "telemetry")]
use rand::{thread_rng, Rng};
use tokio::net::TcpStream;
use tokio::select;
use tokio::sync::{mpsc, watch};
use tokio_util::compat::*;
use futures::io::AsyncReadExt;
use async_trait::async_trait;
use kuska_handshake::async_std::{handshake_server, BoxStream};
use crate::error::*;
use crate::netapp::*;
use crate::proto::*;
use crate::util::*;
// The client and server connection structs (client.rs and server.rs)
// build upon the chunking mechanism which is exclusively contained
// in proto.rs.
// Here, we just care about sending big messages without size limit.
// The format of these messages is described below.
// Chunking happens independently.
// 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
pub(crate) struct ServerConn {
pub(crate) remote_addr: SocketAddr,
pub(crate) peer_id: NodeID,
netapp: Arc<NetApp>,
resp_send: ArcSwapOption<mpsc::UnboundedSender<(RequestID, RequestPriority, Vec<u8>)>>,
}
impl ServerConn {
pub(crate) async fn run(
netapp: Arc<NetApp>,
socket: TcpStream,
must_exit: watch::Receiver<bool>,
) -> Result<(), Error> {
let remote_addr = socket.peer_addr()?;
let mut socket = socket.compat();
let handshake = handshake_server(
&mut socket,
netapp.netid.clone(),
netapp.id,
netapp.privkey.clone(),
)
.await?;
let peer_id = handshake.peer_pk;
debug!(
"Handshake complete (server) with {}@{}",
hex::encode(&peer_id),
remote_addr
);
let (read, write) = socket.split();
let (read, write) =
BoxStream::from_handshake(read, write, handshake, 0x8000).split_read_write();
let (resp_send, resp_recv) = mpsc::unbounded_channel();
let conn = Arc::new(ServerConn {
netapp: netapp.clone(),
remote_addr,
peer_id,
resp_send: ArcSwapOption::new(Some(Arc::new(resp_send))),
});
netapp.connected_as_server(peer_id, conn.clone());
let conn2 = conn.clone();
let recv_future = tokio::spawn(async move {
select! {
r = conn2.recv_loop(read) => r,
_ = await_exit(must_exit) => Ok(())
}
});
let send_future = tokio::spawn(conn.clone().send_loop(resp_recv, write));
recv_future.await.log_err("ServerConn recv_loop");
conn.resp_send.store(None);
send_future.await.log_err("ServerConn send_loop");
netapp.disconnected_as_server(&peer_id, conn);
Ok(())
}
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())?;
2022-02-18 18:01:23 +00:00
let trace_id_len = bytes[2 + path_length] as usize;
let data = &bytes[3 + path_length + trace_id_len..];
let handler_opt = {
let endpoints = self.netapp.endpoints.read().unwrap();
endpoints.get(&path).map(|e| e.clone_endpoint())
};
if let Some(handler) = handler_opt {
2022-02-18 18:01:23 +00:00
cfg_if::cfg_if! {
if #[cfg(feature = "telemetry")] {
use opentelemetry::{
KeyValue,
trace::{FutureExt, TraceContextExt, Tracer},
Context, trace::TraceId
};
let trace_id = if trace_id_len == 16 {
let mut by = [0u8; 16];
by.copy_from_slice(&bytes[3+path_length..19+path_length]);
TraceId::from_bytes(by)
} else {
let mut rng = thread_rng();
TraceId::from_bytes(rng.gen())
};
let tracer = opentelemetry::global::tracer("garage");
let span = tracer
.span_builder(format!("RPC handler {}", path))
.with_trace_id(trace_id)
.with_attributes(vec![
KeyValue::new("path", path),
])
.start(&tracer);
handler.handle(data, self.peer_id)
.with_context(Context::current_with_span(span))
.await
} else {
handler.handle(data, self.peer_id).await
}
}
} else {
Err(Error::NoHandler)
}
}
}
impl SendLoop for ServerConn {}
#[async_trait]
impl RecvLoop for ServerConn {
fn recv_handler(self: &Arc<Self>, id: RequestID, bytes: Vec<u8>) {
let resp_send = self.resp_send.load_full().unwrap();
let self2 = self.clone();
tokio::spawn(async move {
trace!("ServerConn recv_handler {} ({} bytes)", id, bytes.len());
let bytes: Bytes = bytes.into();
let prio = if !bytes.is_empty() { bytes[0] } else { 0u8 };
let resp = self2.recv_handler_aux(&bytes[..]).await;
let mut resp_bytes = vec![];
match resp {
Ok(rb) => {
resp_bytes.push(0u8);
resp_bytes.extend(&rb[..]);
}
Err(e) => {
resp_bytes.push(e.code());
}
}
trace!("ServerConn sending response to {}: ", id);
resp_send
.send((id, prio, resp_bytes))
.log_err("ServerConn recv_handler send resp");
});
}
}