Add basic interceptor support

2019-09-18 23:11:26 -04:00
parent 2ed69f799b
commit e419c18304
5 changed files with 91 additions and 21 deletions
@@ -22,6 +22,7 @@ http = "0.1"
 futures-core-preview = "=0.3.0-alpha.18"
 futures-util-preview = "=0.3.0-alpha.18"
 async-stream = "0.1.1"
+tower = "0.3.0-alpha.1a"

 console = "0.7"
 structopt = "0.2"
@@ -1,6 +1,9 @@
 use structopt::StructOpt;
 use tonic::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)]
 struct Opts {
@@ -26,6 +29,16 @@ async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
        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?;

    Ok(())
@@ -1,20 +1,29 @@
 use super::{
-    service::BoxedIo,
+    service::{BoxedIo, layer_fn},
    tls::{Cert, TlsAcceptor},
 };
 use crate::BoxBody;
 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 hyper::server::{accept::Accept, conn};
 use hyper::Body;
+use std::sync::Arc;
 use std::{
    net::SocketAddr,
    pin::Pin,
    task::{Context, Poll},
+    fmt,
+    future::Future
 };
+use tower::layer::Layer;
 use tower_make::MakeService;
 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)]
 pub struct Server {}
@@ -25,14 +34,15 @@ impl Server {
    }
 }

-#[derive(Debug)]
+#[derive(Default)]
 pub struct Builder {
    tls: Option<(Vec<u8>, Vec<u8>)>,
+    interceptor: Option<Interceptor>,
 }

 impl Builder {
    fn new() -> Self {
-        Self { tls: None }
+        Default::default()
    }

    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 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>
    where
        M: Service<(), Response = S>,
-        M::Error: Into<crate::Error> + 'static,
+        M::Error: Into<crate::Error> + Send + 'static,
        M::Future: Send + 'static,
        S: Service<Request<Body>, Response = Response<BoxBody>> + 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 cert = Cert {
@@ -66,7 +91,10 @@ impl Builder {

        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)
            .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(
    addr: SocketAddr,
    tls: Option<TlsAcceptor>,
@@ -128,40 +162,57 @@ struct Svc<S>(S);
 impl<S> Service<Request<Body>> for Svc<S>
 where
    S: Service<Request<Body>, Response = Response<BoxBody>>,
+    S::Error: Into<crate::Error>
 {
    type Response = Response<BoxBody>;
-    type Error = S::Error;
-    type Future = S::Future;
+    type Error = crate::Error;
+    type Future = MapErr<S::Future, fn(S::Error) -> crate::Error>;

-    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        Ok(()).into()
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.0.poll_ready(cx).map_err(Into::into)
    }

    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>
 where
    M: Service<(), Response = S>,
-    M::Error: Into<crate::Error>,
+    M::Error: Into<crate::Error> + Send,
    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::Error: Into<crate::Error>,
+    S::Error: Into<crate::Error> + Send,
 {
-    type Response = Svc<S>;
-    type Error = M::Error;
-    type Future = MapOk<M::Future, fn(S) -> Svc<S>>;
+    type Response = Either<Svc<S>, BoxService>;
+    type Error = crate::Error;
+    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>> {
-        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 {
-        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)))
+            }
+        })
    }
 }
@@ -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)]
 pub(crate) struct LayerFn<F>(F);

@@ -12,3 +12,4 @@ pub(crate) use self::connection::Connection;
 pub(crate) use self::connector::Connector;
 pub(crate) use self::discover::ServiceList;
 pub(crate) use self::io::BoxedIo;
+pub(crate) use self::layer::layer_fn;