garage/src/rpc/rpc_server.rs

use std::collections::HashMap;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Instant;

use bytes::IntoBuf;
use futures::future::Future;
use futures_util::future::*;
use futures_util::stream::*;
use hyper::server::conn::AddrStream;
use hyper::service::{make_service_fn, service_fn};
use hyper::{Body, Method, Request, Response, Server, StatusCode};
use serde::{Deserialize, Serialize};
use tokio::net::{TcpListener, TcpStream};
use tokio_stream::wrappers::TcpListenerStream;
use tokio_rustls::server::TlsStream;
use tokio_rustls::TlsAcceptor;

use garage_util::config::TlsConfig;
use garage_util::data::*;
use garage_util::error::Error;

use crate::tls_util;

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>,
}

async fn handle_func<M, F, Fut>(
	handler: Arc<F>,
	req: Request<Body>,
	sockaddr: SocketAddr,
	name: Arc<String>,
) -> Result<Response<Body>, Error>
where
	M: RpcMessage + 'static,
	F: Fn(M, SocketAddr) -> Fut + Send + Sync + 'static,
	Fut: Future<Output = Result<M, Error>> + Send + 'static,
{
	let begin_time = Instant::now();
	let whole_body = hyper::body::to_bytes(req.into_body()).await?;
	let msg = rmp_serde::decode::from_read::<_, M>(whole_body.into_buf())?;

	trace!(
		"Request message: {}",
		serde_json::to_string(&msg).unwrap_or("<json error>".into()).chars().take(100).collect::<String>()
	);

	match handler(msg, sockaddr).await {
		Ok(resp) => {
			let resp_bytes = rmp_to_vec_all_named::<Result<M, String>>(&Ok(resp))?;
			let rpc_duration = (Instant::now() - begin_time).as_millis();
			if rpc_duration > 100 {
				debug!("RPC {} ok, took long: {} ms", name, rpc_duration,);
			}
			Ok(Response::new(Body::from(resp_bytes)))
		}
		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() = match e {
				Error::BadRPC(_) => StatusCode::BAD_REQUEST,
				_ => StatusCode::INTERNAL_SERVER_ERROR,
			};
			warn!(
				"RPC error ({}): {} ({} ms)",
				name,
				e,
				(Instant::now() - begin_time).as_millis(),
			);
			Ok(err_response)
		}
	}
}

impl RpcServer {
	pub fn new(bind_addr: SocketAddr, tls_config: Option<TlsConfig>) -> Self {
		Self {
			bind_addr,
			tls_config,
			handlers: HashMap::new(),
		}
	}

	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,
	{
		let name2 = Arc::new(name.clone());
		let handler_arc = Arc::new(handler);
		let handler = Box::new(move |req: Request<Body>, sockaddr: SocketAddr| {
			let handler2 = handler_arc.clone();
			let b: ResponseFuture = Box::pin(handle_func(handler2, req, sockaddr, name2.clone()));
			b
		});
		self.handlers.insert(name, handler);
	}

	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);
		}

		let path = &req.uri().path()[1..].to_string();

		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);
			}
		};

		trace!("({}) Handling request", path);

		let resp_waiter = tokio::spawn(handler(req, addr));
		match resp_waiter.await {
			Err(err) => {
				warn!("Handler await error: {}", err);
				let mut ise = Response::default();
				*ise.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
				Ok(ise)
			}
			Ok(Err(err)) => {
				trace!("({}) Request handler failed: {}", path, err);
				let mut bad_request = Response::new(Body::from(format!("{}", err)));
				*bad_request.status_mut() = StatusCode::BAD_REQUEST;
				Ok(bad_request)
			}
			Ok(Ok(resp)) => {
				trace!("({}) Request handler succeeded", path);
				Ok(resp)
			}
		}
	}

	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)?;
			}

			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 listener = TcpListener::bind(&self.bind_addr).await?;
			let incoming = TcpListenerStream::new(listener).filter_map(|socket| async {
				match socket {
					Ok(stream) => match tls_acceptor.clone().accept(stream).await {
						Ok(x) => Some(Ok::<_, hyper::Error>(x)),
						Err(_e) => None,
					},
					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| {
							warn!("RPC handler error: {}", e);
							e
						})
					}))
				}
			});

			let server = Server::builder(incoming).serve(service);

			let graceful = server.with_graceful_shutdown(shutdown_signal);
			info!("RPC server listening on http://{}", self.bind_addr);

			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| {
							warn!("RPC handler error: {}", e);
							e
						})
					}))
				}
			});

			let server = Server::bind(&self.bind_addr).serve(service);

			let graceful = server.with_graceful_shutdown(shutdown_signal);
			info!("RPC server listening on http://{}", self.bind_addr);

			graceful.await?;
		}

		Ok(())
	}
}
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`use std::collections::HashMap;`
Some more basic work 2020-04-06 17:55:39 +00:00			`use std::net::SocketAddr;`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`use std::pin::Pin;`
Some more basic work 2020-04-06 17:55:39 +00:00			`use std::sync::Arc;`
Add node tags in configuration 2020-04-21 14:07:15 +00:00			`use std::time::Instant;`
Some more basic work 2020-04-06 17:55:39 +00:00
First commit: skeleton for something great 2020-04-05 21:33:42 +00:00			`use bytes::IntoBuf;`
Fix the Sync issue. Details: So the HTTP client future of Hyper is not Sync, thus the stream that read blocks wasn't either. However Hyper's default Body type requires a stream to be Sync for wrap_stream. Solution: reimplement a custom HTTP body type. 2020-04-10 20:01:48 +00:00			`use futures::future::Future;`
We don't want the debugging to fail, actually. 2020-04-10 21:43:35 +00:00			`use futures_util::future::*;`
Trying to do TLS 2020-04-12 13:51:19 +00:00			`use futures_util::stream::*;`
Some more basic work 2020-04-06 17:55:39 +00:00			`use hyper::server::conn::AddrStream;`
Fix the Sync issue. Details: So the HTTP client future of Hyper is not Sync, thus the stream that read blocks wasn't either. However Hyper's default Body type requires a stream to be Sync for wrap_stream. Solution: reimplement a custom HTTP body type. 2020-04-10 20:01:48 +00:00			`use hyper::service::{make_service_fn, service_fn};`
First commit: skeleton for something great 2020-04-05 21:33:42 +00:00			`use hyper::{Body, Method, Request, Response, Server, StatusCode};`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`use serde::{Deserialize, Serialize};`
Trying to do TLS 2020-04-12 13:51:19 +00:00			`use tokio::net::{TcpListener, TcpStream};`
WIP migrate to tokio 1 2021-03-15 21:36:41 +00:00			`use tokio_stream::wrappers::TcpListenerStream;`
Trying to do TLS 2020-04-12 13:51:19 +00:00			`use tokio_rustls::server::TlsStream;`
			`use tokio_rustls::TlsAcceptor;`
First commit: skeleton for something great 2020-04-05 21:33:42 +00:00
Split code for modular compilation 2020-04-24 10:10:01 +00:00			`use garage_util::config::TlsConfig;`
			`use garage_util::data::*;`
			`use garage_util::error::Error;`
Reorganize code 2020-04-23 17:05:46 +00:00
Split code for modular compilation 2020-04-24 10:10:01 +00:00			`use crate::tls_util;`
Why is it not Sync?? 2020-04-09 21:45:07 +00:00
Massive RPC refactoring 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>,`
Some more basic work 2020-04-06 17:55:39 +00:00			`}`

Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`async fn handle_func<M, F, Fut>(`
			`handler: Arc<F>,`
Fix the Sync issue. Details: So the HTTP client future of Hyper is not Sync, thus the stream that read blocks wasn't either. However Hyper's default Body type requires a stream to be Sync for wrap_stream. Solution: reimplement a custom HTTP body type. 2020-04-10 20:01:48 +00:00			`req: Request<Body>,`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`sockaddr: SocketAddr,`
Pretty logging 2020-04-21 12:54:55 +00:00			`name: Arc<String>,`
Massive RPC refactoring 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,`
			`{`
Pretty logging 2020-04-21 12:54:55 +00:00			`let begin_time = Instant::now();`
First commit: skeleton for something great 2020-04-05 21:33:42 +00:00			`let whole_body = hyper::body::to_bytes(req.into_body()).await?;`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`let msg = rmp_serde::decode::from_read::<_, M>(whole_body.into_buf())?;`
Update sled & try to debug deadlock (but its in sled...) 2021-02-23 20:27:28 +00:00
			`trace!(`
			`"Request message: {}",`
Tiny things 2021-03-15 17:40:27 +00:00			`serde_json::to_string(&msg).unwrap_or("<json error>".into()).chars().take(100).collect::<String>()`
Update sled & try to debug deadlock (but its in sled...) 2021-02-23 20:27:28 +00:00			`);`

Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`match handler(msg, sockaddr).await {`
			`Ok(resp) => {`
			`let resp_bytes = rmp_to_vec_all_named::<Result<M, String>>(&Ok(resp))?;`
Well they still have to exit when we're exiting though 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,);`
			`}`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`Ok(Response::new(Body::from(resp_bytes)))`
fixes - make block_put call uninterruptible by client - used meta_replication_factor instead of data_replication_factor - listen on ipv6 2020-04-12 10:10:33 +00:00			`}`
Massive RPC refactoring 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));`
Refactor error management in API part 2020-11-08 14:04:30 +00:00			`*err_response.status_mut() = match e {`
			`Error::BadRPC(_) => StatusCode::BAD_REQUEST,`
			`_ => StatusCode::INTERNAL_SERVER_ERROR,`
			`};`
Add node tags in configuration 2020-04-21 14:07:15 +00:00			`warn!(`
Well they still have to exit when we're exiting though 2020-04-22 17:04:33 +00:00			`"RPC error ({}): {} ({} ms)",`
Add node tags in configuration 2020-04-21 14:07:15 +00:00			`name,`
			`e,`
			`(Instant::now() - begin_time).as_millis(),`
			`);`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`Ok(err_response)`
Some work in actually storing things 2020-04-08 20:00:41 +00:00			`}`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`}`
First commit: skeleton for something great 2020-04-05 21:33:42 +00:00			`}`

Massive RPC refactoring 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(),`
			`}`
			`}`
First commit: skeleton for something great 2020-04-05 21:33:42 +00:00
Massive RPC refactoring 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,`
			`{`
Pretty logging 2020-04-21 12:54:55 +00:00			`let name2 = Arc::new(name.clone());`
Massive RPC refactoring 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();`
Pretty logging 2020-04-21 12:54:55 +00:00			`let b: ResponseFuture = Box::pin(handle_func(handler2, req, sockaddr, name2.clone()));`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`b`
			`});`
			`self.handlers.insert(name, handler);`
			`}`
Trying to do TLS 2020-04-12 13:51:19 +00:00
Massive RPC refactoring 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);`
Some more basic work 2020-04-06 17:55:39 +00:00			`}`
First commit: skeleton for something great 2020-04-05 21:33:42 +00:00
Update sled & try to debug deadlock (but its in sled...) 2021-02-23 20:27:28 +00:00			`let path = &req.uri().path()[1..].to_string();`

Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`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);`
Trying to do TLS 2020-04-12 13:51:19 +00:00			`}`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`};`

Update sled & try to debug deadlock (but its in sled...) 2021-02-23 20:27:28 +00:00			`trace!("({}) Handling request", path);`

Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`let resp_waiter = tokio::spawn(handler(req, addr));`
			`match resp_waiter.await {`
Begin implement bucket management & admin commands 2020-04-19 15:15:48 +00:00			`Err(err) => {`
Pretty logging 2020-04-21 12:54:55 +00:00			`warn!("Handler await error: {}", err);`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`let mut ise = Response::default();`
			`*ise.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;`
			`Ok(ise)`
Trying to do TLS 2020-04-12 13:51:19 +00:00			`}`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`Ok(Err(err)) => {`
Update sled & try to debug deadlock (but its in sled...) 2021-02-23 20:27:28 +00:00			`trace!("({}) Request handler failed: {}", path, err);`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`let mut bad_request = Response::new(Body::from(format!("{}", err)));`
			`*bad_request.status_mut() = StatusCode::BAD_REQUEST;`
			`Ok(bad_request)`
			`}`
Update sled & try to debug deadlock (but its in sled...) 2021-02-23 20:27:28 +00:00			`Ok(Ok(resp)) => {`
			`trace!("({}) Request handler succeeded", path);`
			`Ok(resp)`
			`}`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`}`
			`}`
Trying to do TLS 2020-04-12 13:51:19 +00:00
Massive RPC refactoring 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)?;`
Trying to do TLS 2020-04-12 13:51:19 +00:00			`}`

Massive RPC refactoring 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)));`

WIP migrate to tokio 1 2021-03-15 21:36:41 +00:00			`let listener = TcpListener::bind(&self.bind_addr).await?;`
			`let incoming = TcpListenerStream::new(listener).filter_map(\|socket\| async {`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`match socket {`
			`Ok(stream) => match tls_acceptor.clone().accept(stream).await {`
			`Ok(x) => Some(Ok::<_, hyper::Error>(x)),`
More aggressive sync timings & improve other stuff 2020-04-19 17:59:59 +00:00			`Err(_e) => None,`
Massive RPC refactoring 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\| {`
Pretty logging 2020-04-21 12:54:55 +00:00			`warn!("RPC handler error: {}", e);`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`e`
			`})`
			`}))`
			`}`
			`});`

			`let server = Server::builder(incoming).serve(service);`

			`let graceful = server.with_graceful_shutdown(shutdown_signal);`
Pretty logging 2020-04-21 12:54:55 +00:00			`info!("RPC server listening on http://{}", self.bind_addr);`
Massive RPC refactoring 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\| {`
Pretty logging 2020-04-21 12:54:55 +00:00			`warn!("RPC handler error: {}", e);`
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`e`
			`})`
			`}))`
			`}`
			`});`
First commit: skeleton for something great 2020-04-05 21:33:42 +00:00
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`let server = Server::bind(&self.bind_addr).serve(service);`
Trying to do TLS 2020-04-12 13:51:19 +00:00
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`let graceful = server.with_graceful_shutdown(shutdown_signal);`
Pretty logging 2020-04-21 12:54:55 +00:00			`info!("RPC server listening on http://{}", self.bind_addr);`
First commit: skeleton for something great 2020-04-05 21:33:42 +00:00
Massive RPC refactoring 2020-04-18 17:21:34 +00:00			`graceful.await?;`
			`}`

			`Ok(())`
			`}`
First commit: skeleton for something great 2020-04-05 21:33:42 +00:00			`}`