talksik/tonic
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add basic interceptor support

Browse Source
This commit is contained in:
Lucio Franco
2019-09-18 23:16:48 -04:00
parent 2ed69f799b
commit e419c18304
5 changed files with 91 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tonic-interop/Cargo.toml
+1
View File
@@ -22,6 +22,7 @@ http = "0.1"
futures-core-preview = "=0.3.0-alpha.18" futures-core-preview = "=0.3.0-alpha.18"
futures-util-preview = "=0.3.0-alpha.18" futures-util-preview = "=0.3.0-alpha.18"
async-stream = "0.1.1" async-stream = "0.1.1"
tower = "0.3.0-alpha.1a"
console = "0.7" console = "0.7"
structopt = "0.2" structopt = "0.2"
tonic-interop/src/bin/server.rs
+13
View File
@@ -1,6 +1,9 @@
use structopt::StructOpt; use structopt::StructOpt;
use tonic::Server; use tonic::Server;
use tonic_interop::server; use tonic_interop::server;
// TODO: move GrpcService out of client since it can be used for the
// server too.
use tonic::client::GrpcService;
#[derive(StructOpt)] #[derive(StructOpt)]
struct Opts { struct Opts {
@@ -26,6 +29,16 @@ async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
builder.tls(ca, key); builder.tls(ca, key);
} }
builder.interceptor_fn(|svc, req| {
println!("INBOUND REQUEST={:?}", req);
let call = svc.call(req);
async move {
let res = call.await?;
println!("OUTBOUND RESPONSE={:?}", res);
Ok(res)
}
});
builder.serve(addr, test_service).await?; builder.serve(addr, test_service).await?;
Ok(()) Ok(())
tonic/src/transport/server.rs
+72 -21
View File
@@ -1,20 +1,29 @@
use super::{ use super::{
service::BoxedIo, service::{BoxedIo, layer_fn},
tls::{Cert, TlsAcceptor}, tls::{Cert, TlsAcceptor},
}; };
use crate::BoxBody; use crate::BoxBody;
use futures_core::Stream; use futures_core::Stream;
use futures_util::{ready, try_future::MapOk, TryFutureExt, TryStreamExt}; use futures_util::{ready, try_future::MapErr, TryFutureExt, TryStreamExt};
use http::{Request, Response}; use http::{Request, Response};
use hyper::server::{accept::Accept, conn}; use hyper::server::{accept::Accept, conn};
use hyper::Body; use hyper::Body;
use std::sync::Arc;
use std::{ use std::{
net::SocketAddr, net::SocketAddr,
pin::Pin, pin::Pin,
task::{Context, Poll}, task::{Context, Poll},
fmt,
future::Future
}; };
use tower::layer::Layer;
use tower_make::MakeService; use tower_make::MakeService;
use tower_service::Service; use tower_service::Service;
use tower::layer::util::Stack;
use tower::util::Either;
type BoxService = tower::util::BoxService<Request<Body>, Response<BoxBody>, crate::Error>;
type Interceptor = Arc<dyn Layer<BoxService, Service = BoxService> + Send + Sync + 'static>;
#[derive(Debug)] #[derive(Debug)]
pub struct Server {} pub struct Server {}
@@ -25,14 +34,15 @@ impl Server {
} }
} }
#[derive(Debug)] #[derive(Default)]
pub struct Builder { pub struct Builder {
tls: Option<(Vec<u8>, Vec<u8>)>, tls: Option<(Vec<u8>, Vec<u8>)>,
interceptor: Option<Interceptor>,
} }
impl Builder { impl Builder {
fn new() -> Self { fn new() -> Self {
Self { tls: None } Default::default()
} }
pub fn tls(&mut self, pem: Vec<u8>, key: Vec<u8>) -> &mut Self { pub fn tls(&mut self, pem: Vec<u8>, key: Vec<u8>) -> &mut Self {
@@ -43,14 +53,29 @@ impl Builder {
// pub fn concurrency_limit(&mut self, limit: usize) -> &mut Self { // pub fn concurrency_limit(&mut self, limit: usize) -> &mut Self {
// } // }
pub fn interceptor_fn<F, Out>(&mut self, f: F) -> &mut Self
where
F: Fn(&mut BoxService, Request<Body>) -> Out + Send + Sync + 'static,
Out: Future<Output = Result<Response<BoxBody>, crate::Error>> + Send + 'static
{
let f = Arc::new(f);
let interceptor = layer_fn(move |mut s| {
let f = f.clone();
tower::service_fn(move |req| f(&mut s, req))
});
let layer = Stack::new(interceptor, layer_fn(|s| BoxService::new(s)));
self.interceptor = Some(Arc::new(layer));
self
}
pub async fn serve<M, S>(self, addr: SocketAddr, svc: M) -> Result<(), super::Error> pub async fn serve<M, S>(self, addr: SocketAddr, svc: M) -> Result<(), super::Error>
where where
M: Service<(), Response = S>, M: Service<(), Response = S>,
M::Error: Into<crate::Error> + 'static, M::Error: Into<crate::Error> + Send + 'static,
M::Future: Send + 'static, M::Future: Send + 'static,
S: Service<Request<Body>, Response = Response<BoxBody>> + Send + 'static, S: Service<Request<Body>, Response = Response<BoxBody>> + Send + 'static,
S::Future: Send + 'static, S::Future: Send + 'static,
S::Error: Into<crate::Error>, S::Error: Into<crate::Error> + Send,
{ {
let tls = if let Some(tls) = self.tls { let tls = if let Some(tls) = self.tls {
let cert = Cert { let cert = Cert {
@@ -66,7 +91,10 @@ impl Builder {
let incoming = hyper::server::accept::from_stream(incoming(addr, tls)); let incoming = hyper::server::accept::from_stream(incoming(addr, tls));
let svc = MakeSvc(svc); let svc = MakeSvc {
inner: svc,
interceptor: self.interceptor.clone(),
};
hyper::Server::builder(incoming) hyper::Server::builder(incoming)
.http2_only(true) .http2_only(true)
@@ -78,6 +106,12 @@ impl Builder {
} }
} }
impl fmt::Debug for Builder {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Builder").finish()
}
}
fn incoming( fn incoming(
addr: SocketAddr, addr: SocketAddr,
tls: Option<TlsAcceptor>, tls: Option<TlsAcceptor>,
@@ -128,40 +162,57 @@ struct Svc<S>(S);
impl<S> Service<Request<Body>> for Svc<S> impl<S> Service<Request<Body>> for Svc<S>
where where
S: Service<Request<Body>, Response = Response<BoxBody>>, S: Service<Request<Body>, Response = Response<BoxBody>>,
S::Error: Into<crate::Error>
{ {
type Response = Response<BoxBody>; type Response = Response<BoxBody>;
type Error = S::Error; type Error = crate::Error;
type Future = S::Future; type Future = MapErr<S::Future, fn(S::Error) -> crate::Error>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Ok(()).into() self.0.poll_ready(cx).map_err(Into::into)
} }
fn call(&mut self, req: Request<Body>) -> Self::Future { fn call(&mut self, req: Request<Body>) -> Self::Future {
self.0.call(req) self.0.call(req).map_err(|e| e.into())
} }
} }
struct MakeSvc<M>(M); struct MakeSvc<M> {
interceptor: Option<Interceptor>,
inner: M,
}
impl<M, S, T> Service<T> for MakeSvc<M> impl<M, S, T> Service<T> for MakeSvc<M>
where where
M: Service<(), Response = S>, M: Service<(), Response = S>,
M::Error: Into<crate::Error>, M::Error: Into<crate::Error> + Send,
M::Future: Send + 'static, M::Future: Send + 'static,
S: Service<Request<Body>, Response = Response<BoxBody>>, S: Service<Request<Body>, Response = Response<BoxBody>> + Send + 'static,
S::Future: Send + 'static, S::Future: Send + 'static,
S::Error: Into<crate::Error>, S::Error: Into<crate::Error> + Send,
{ {
type Response = Svc<S>; type Response = Either<Svc<S>, BoxService>;
type Error = M::Error; type Error = crate::Error;
type Future = MapOk<M::Future, fn(S) -> Svc<S>>; type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
MakeService::poll_ready(&mut self.0, cx) MakeService::poll_ready(&mut self.inner, cx).map_err(Into::into)
} }
fn call(&mut self, _: T) -> Self::Future { fn call(&mut self, _: T) -> Self::Future {
self.0.make_service(()).map_ok(|s| Svc(s)) // self.inner.make_service(()).map_ok(|s| Svc(s))
let interceptor = self.interceptor.clone();
// self.inner.make_service(()).map_ok(|s| intercept.layer(BoxService::new(Svc(s))))
let make = self.inner.make_service(());
Box::pin(async move {
let svc = make.await.map_err(Into::into)?;
if let Some(interceptor) = interceptor {
let layered = interceptor.layer(BoxService::new(Svc(svc)));
Ok(Either::B(layered))
} else {
Ok(Either::A(Svc(svc)))
}
})
} }
} }
tonic/src/transport/service/layer.rs
+4
View File
@@ -41,6 +41,10 @@ impl<L> ServiceBuilderExt<L> for ServiceBuilder<L> {
} }
} }
pub(crate) fn layer_fn<F>(f: F) -> LayerFn<F> {
LayerFn(f)
}
#[derive(Clone, Copy, Debug)] #[derive(Clone, Copy, Debug)]
pub(crate) struct LayerFn<F>(F); pub(crate) struct LayerFn<F>(F);
tonic/src/transport/service/mod.rs
+1
View File
@@ -12,3 +12,4 @@ pub(crate) use self::connection::Connection;
pub(crate) use self::connector::Connector; pub(crate) use self::connector::Connector;
pub(crate) use self::discover::ServiceList; pub(crate) use self::discover::ServiceList;
pub(crate) use self::io::BoxedIo; pub(crate) use self::io::BoxedIo;
pub(crate) use self::layer::layer_fn;