2020-04-18 17:21:34 +00:00
|
|
|
use std::collections::HashMap;
|
2020-04-06 17:55:39 +00:00
|
|
|
use std::net::SocketAddr;
|
2020-04-18 17:21:34 +00:00
|
|
|
use std::pin::Pin;
|
2020-04-06 17:55:39 +00:00
|
|
|
use std::sync::Arc;
|
2020-04-21 14:07:15 +00:00
|
|
|
use std::time::Instant;
|
2020-04-06 17:55:39 +00:00
|
|
|
|
2020-04-05 21:33:42 +00:00
|
|
|
use bytes::IntoBuf;
|
2020-04-10 20:01:48 +00:00
|
|
|
use futures::future::Future;
|
2020-04-10 21:43:35 +00:00
|
|
|
use futures_util::future::*;
|
2020-04-12 13:51:19 +00:00
|
|
|
use futures_util::stream::*;
|
2020-04-06 17:55:39 +00:00
|
|
|
use hyper::server::conn::AddrStream;
|
2020-04-10 20:01:48 +00:00
|
|
|
use hyper::service::{make_service_fn, service_fn};
|
2020-04-05 21:33:42 +00:00
|
|
|
use hyper::{Body, Method, Request, Response, Server, StatusCode};
|
2020-04-18 17:21:34 +00:00
|
|
|
use serde::{Deserialize, Serialize};
|
2020-04-12 13:51:19 +00:00
|
|
|
use tokio::net::{TcpListener, TcpStream};
|
|
|
|
|
use tokio_rustls::server::TlsStream;
|
|
|
|
|
use tokio_rustls::TlsAcceptor;
|
2020-04-05 21:33:42 +00:00
|
|
|
|
2020-04-23 17:05:46 +00:00
|
|
|
use crate::config::TlsConfig;
|
2020-04-16 21:13:15 +00:00
|
|
|
use crate::data::*;
|
2020-04-10 20:01:48 +00:00
|
|
|
use crate::error::Error;
|
2020-04-23 17:05:46 +00:00
|
|
|
|
|
|
|
|
use crate::rpc::tls_util;
|
2020-04-09 21:45:07 +00:00
|
|
|
|
2020-04-18 17:21:34 +00:00
|
|
|
pub trait RpcMessage: Serialize + for<'de> Deserialize<'de> + Send + Sync {}
|
|
|
|
|
|
|
|
|
|
type ResponseFuture = Pin<Box<dyn Future<Output = Result<Response<Body>, Error>> + Send>>;
|
|
|
|
|
type Handler = Box<dyn Fn(Request<Body>, SocketAddr) -> ResponseFuture + Send + Sync>;
|
|
|
|
|
|
|
|
|
|
pub struct RpcServer {
|
|
|
|
|
pub bind_addr: SocketAddr,
|
|
|
|
|
pub tls_config: Option<TlsConfig>,
|
|
|
|
|
|
|
|
|
|
handlers: HashMap<String, Handler>,
|
2020-04-06 17:55:39 +00:00
|
|
|
}
|
|
|
|
|
|
2020-04-18 17:21:34 +00:00
|
|
|
async fn handle_func<M, F, Fut>(
|
|
|
|
|
handler: Arc<F>,
|
2020-04-10 20:01:48 +00:00
|
|
|
req: Request<Body>,
|
2020-04-18 17:21:34 +00:00
|
|
|
sockaddr: SocketAddr,
|
2020-04-21 12:54:55 +00:00
|
|
|
name: Arc<String>,
|
2020-04-18 17:21:34 +00:00
|
|
|
) -> Result<Response<Body>, Error>
|
|
|
|
|
where
|
|
|
|
|
M: RpcMessage + 'static,
|
|
|
|
|
F: Fn(M, SocketAddr) -> Fut + Send + Sync + 'static,
|
|
|
|
|
Fut: Future<Output = Result<M, Error>> + Send + 'static,
|
|
|
|
|
{
|
2020-04-21 12:54:55 +00:00
|
|
|
let begin_time = Instant::now();
|
2020-04-05 21:33:42 +00:00
|
|
|
let whole_body = hyper::body::to_bytes(req.into_body()).await?;
|
2020-04-18 17:21:34 +00:00
|
|
|
let msg = rmp_serde::decode::from_read::<_, M>(whole_body.into_buf())?;
|
|
|
|
|
match handler(msg, sockaddr).await {
|
|
|
|
|
Ok(resp) => {
|
|
|
|
|
let resp_bytes = rmp_to_vec_all_named::<Result<M, String>>(&Ok(resp))?;
|
2020-04-22 17:04:33 +00:00
|
|
|
let rpc_duration = (Instant::now() - begin_time).as_millis();
|
|
|
|
|
if rpc_duration > 100 {
|
|
|
|
|
debug!("RPC {} ok, took long: {} ms", name, rpc_duration,);
|
|
|
|
|
}
|
2020-04-18 17:21:34 +00:00
|
|
|
Ok(Response::new(Body::from(resp_bytes)))
|
2020-04-12 10:10:33 +00:00
|
|
|
}
|
2020-04-18 17:21:34 +00:00
|
|
|
Err(e) => {
|
|
|
|
|
let err_str = format!("{}", e);
|
|
|
|
|
let rep_bytes = rmp_to_vec_all_named::<Result<M, String>>(&Err(err_str))?;
|
|
|
|
|
let mut err_response = Response::new(Body::from(rep_bytes));
|
|
|
|
|
*err_response.status_mut() = e.http_status_code();
|
2020-04-21 14:07:15 +00:00
|
|
|
warn!(
|
2020-04-22 17:04:33 +00:00
|
|
|
"RPC error ({}): {} ({} ms)",
|
2020-04-21 14:07:15 +00:00
|
|
|
name,
|
|
|
|
|
e,
|
|
|
|
|
(Instant::now() - begin_time).as_millis(),
|
|
|
|
|
);
|
2020-04-18 17:21:34 +00:00
|
|
|
Ok(err_response)
|
2020-04-08 20:00:41 +00:00
|
|
|
}
|
2020-04-18 17:21:34 +00:00
|
|
|
}
|
2020-04-05 21:33:42 +00:00
|
|
|
}
|
|
|
|
|
|
2020-04-18 17:21:34 +00:00
|
|
|
impl RpcServer {
|
|
|
|
|
pub fn new(bind_addr: SocketAddr, tls_config: Option<TlsConfig>) -> Self {
|
|
|
|
|
Self {
|
|
|
|
|
bind_addr,
|
|
|
|
|
tls_config,
|
|
|
|
|
handlers: HashMap::new(),
|
|
|
|
|
}
|
|
|
|
|
}
|
2020-04-05 21:33:42 +00:00
|
|
|
|
2020-04-18 17:21:34 +00:00
|
|
|
pub fn add_handler<M, F, Fut>(&mut self, name: String, handler: F)
|
|
|
|
|
where
|
|
|
|
|
M: RpcMessage + 'static,
|
|
|
|
|
F: Fn(M, SocketAddr) -> Fut + Send + Sync + 'static,
|
|
|
|
|
Fut: Future<Output = Result<M, Error>> + Send + 'static,
|
|
|
|
|
{
|
2020-04-21 12:54:55 +00:00
|
|
|
let name2 = Arc::new(name.clone());
|
2020-04-18 17:21:34 +00:00
|
|
|
let handler_arc = Arc::new(handler);
|
|
|
|
|
let handler = Box::new(move |req: Request<Body>, sockaddr: SocketAddr| {
|
|
|
|
|
let handler2 = handler_arc.clone();
|
2020-04-21 12:54:55 +00:00
|
|
|
let b: ResponseFuture = Box::pin(handle_func(handler2, req, sockaddr, name2.clone()));
|
2020-04-18 17:21:34 +00:00
|
|
|
b
|
|
|
|
|
});
|
|
|
|
|
self.handlers.insert(name, handler);
|
|
|
|
|
}
|
2020-04-12 13:51:19 +00:00
|
|
|
|
2020-04-18 17:21:34 +00:00
|
|
|
async fn handler(
|
|
|
|
|
self: Arc<Self>,
|
|
|
|
|
req: Request<Body>,
|
|
|
|
|
addr: SocketAddr,
|
|
|
|
|
) -> Result<Response<Body>, Error> {
|
|
|
|
|
if req.method() != &Method::POST {
|
|
|
|
|
let mut bad_request = Response::default();
|
|
|
|
|
*bad_request.status_mut() = StatusCode::BAD_REQUEST;
|
|
|
|
|
return Ok(bad_request);
|
2020-04-06 17:55:39 +00:00
|
|
|
}
|
2020-04-05 21:33:42 +00:00
|
|
|
|
2020-04-18 17:21:34 +00:00
|
|
|
let path = &req.uri().path()[1..];
|
|
|
|
|
let handler = match self.handlers.get(path) {
|
|
|
|
|
Some(h) => h,
|
|
|
|
|
None => {
|
|
|
|
|
let mut not_found = Response::default();
|
|
|
|
|
*not_found.status_mut() = StatusCode::NOT_FOUND;
|
|
|
|
|
return Ok(not_found);
|
2020-04-12 13:51:19 +00:00
|
|
|
}
|
2020-04-18 17:21:34 +00:00
|
|
|
};
|
|
|
|
|
|
|
|
|
|
let resp_waiter = tokio::spawn(handler(req, addr));
|
|
|
|
|
match resp_waiter.await {
|
2020-04-19 15:15:48 +00:00
|
|
|
Err(err) => {
|
2020-04-21 12:54:55 +00:00
|
|
|
warn!("Handler await error: {}", err);
|
2020-04-18 17:21:34 +00:00
|
|
|
let mut ise = Response::default();
|
|
|
|
|
*ise.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
|
|
|
|
|
Ok(ise)
|
2020-04-12 13:51:19 +00:00
|
|
|
}
|
2020-04-18 17:21:34 +00:00
|
|
|
Ok(Err(err)) => {
|
|
|
|
|
let mut bad_request = Response::new(Body::from(format!("{}", err)));
|
|
|
|
|
*bad_request.status_mut() = StatusCode::BAD_REQUEST;
|
|
|
|
|
Ok(bad_request)
|
|
|
|
|
}
|
|
|
|
|
Ok(Ok(resp)) => Ok(resp),
|
|
|
|
|
}
|
|
|
|
|
}
|
2020-04-12 13:51:19 +00:00
|
|
|
|
2020-04-18 17:21:34 +00:00
|
|
|
pub async fn run(
|
|
|
|
|
self: Arc<Self>,
|
|
|
|
|
shutdown_signal: impl Future<Output = ()>,
|
|
|
|
|
) -> Result<(), Error> {
|
|
|
|
|
if let Some(tls_config) = self.tls_config.as_ref() {
|
|
|
|
|
let ca_certs = tls_util::load_certs(&tls_config.ca_cert)?;
|
|
|
|
|
let node_certs = tls_util::load_certs(&tls_config.node_cert)?;
|
|
|
|
|
let node_key = tls_util::load_private_key(&tls_config.node_key)?;
|
|
|
|
|
|
|
|
|
|
let mut ca_store = rustls::RootCertStore::empty();
|
|
|
|
|
for crt in ca_certs.iter() {
|
|
|
|
|
ca_store.add(crt)?;
|
2020-04-12 13:51:19 +00:00
|
|
|
}
|
|
|
|
|
|
2020-04-18 17:21:34 +00:00
|
|
|
let mut config =
|
|
|
|
|
rustls::ServerConfig::new(rustls::AllowAnyAuthenticatedClient::new(ca_store));
|
|
|
|
|
config.set_single_cert([&node_certs[..], &ca_certs[..]].concat(), node_key)?;
|
|
|
|
|
let tls_acceptor = Arc::new(TlsAcceptor::from(Arc::new(config)));
|
|
|
|
|
|
|
|
|
|
let mut listener = TcpListener::bind(&self.bind_addr).await?;
|
|
|
|
|
let incoming = listener.incoming().filter_map(|socket| async {
|
|
|
|
|
match socket {
|
|
|
|
|
Ok(stream) => match tls_acceptor.clone().accept(stream).await {
|
|
|
|
|
Ok(x) => Some(Ok::<_, hyper::Error>(x)),
|
2020-04-19 17:59:59 +00:00
|
|
|
Err(_e) => None,
|
2020-04-18 17:21:34 +00:00
|
|
|
},
|
|
|
|
|
Err(_) => None,
|
|
|
|
|
}
|
|
|
|
|
});
|
|
|
|
|
let incoming = hyper::server::accept::from_stream(incoming);
|
|
|
|
|
|
|
|
|
|
let self_arc = self.clone();
|
|
|
|
|
let service = make_service_fn(|conn: &TlsStream<TcpStream>| {
|
|
|
|
|
let client_addr = conn
|
|
|
|
|
.get_ref()
|
|
|
|
|
.0
|
|
|
|
|
.peer_addr()
|
|
|
|
|
.unwrap_or(([0, 0, 0, 0], 0).into());
|
|
|
|
|
let self_arc = self_arc.clone();
|
|
|
|
|
async move {
|
|
|
|
|
Ok::<_, Error>(service_fn(move |req: Request<Body>| {
|
|
|
|
|
self_arc.clone().handler(req, client_addr).map_err(|e| {
|
2020-04-21 12:54:55 +00:00
|
|
|
warn!("RPC handler error: {}", e);
|
2020-04-18 17:21:34 +00:00
|
|
|
e
|
|
|
|
|
})
|
|
|
|
|
}))
|
|
|
|
|
}
|
|
|
|
|
});
|
|
|
|
|
|
|
|
|
|
let server = Server::builder(incoming).serve(service);
|
|
|
|
|
|
|
|
|
|
let graceful = server.with_graceful_shutdown(shutdown_signal);
|
2020-04-21 12:54:55 +00:00
|
|
|
info!("RPC server listening on http://{}", self.bind_addr);
|
2020-04-18 17:21:34 +00:00
|
|
|
|
|
|
|
|
graceful.await?;
|
|
|
|
|
} else {
|
|
|
|
|
let self_arc = self.clone();
|
|
|
|
|
let service = make_service_fn(move |conn: &AddrStream| {
|
|
|
|
|
let client_addr = conn.remote_addr();
|
|
|
|
|
let self_arc = self_arc.clone();
|
|
|
|
|
async move {
|
|
|
|
|
Ok::<_, Error>(service_fn(move |req: Request<Body>| {
|
|
|
|
|
self_arc.clone().handler(req, client_addr).map_err(|e| {
|
2020-04-21 12:54:55 +00:00
|
|
|
warn!("RPC handler error: {}", e);
|
2020-04-18 17:21:34 +00:00
|
|
|
e
|
|
|
|
|
})
|
|
|
|
|
}))
|
|
|
|
|
}
|
|
|
|
|
});
|
2020-04-05 21:33:42 +00:00
|
|
|
|
2020-04-18 17:21:34 +00:00
|
|
|
let server = Server::bind(&self.bind_addr).serve(service);
|
2020-04-12 13:51:19 +00:00
|
|
|
|
2020-04-18 17:21:34 +00:00
|
|
|
let graceful = server.with_graceful_shutdown(shutdown_signal);
|
2020-04-21 12:54:55 +00:00
|
|
|
info!("RPC server listening on http://{}", self.bind_addr);
|
2020-04-05 21:33:42 +00:00
|
|
|
|
2020-04-18 17:21:34 +00:00
|
|
|
graceful.await?;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
|
}
|
2020-04-05 21:33:42 +00:00
|
|
|
}
|
