feat: Add gRPC interceptors (#232)

This change introduces proper gRPC interceptors that are avilable regardless of the transport used. Each codegen service now produces an additional method called `with_interceptor` that accepts a `Interceptor`. All examples have been updated to use this new style and interop has a custom `tower::Service` middleware to echo the headers. There is also a new `interceptor` example that shows basic usage. BREAKING CHANGE: removed `interceptor_fn` and `intercep_headers_fn` from `transport` in favor of using `tonic::Interceptor`.
2020-01-13 21:26:08 -05:00
parent 0fa2bf1cea
commit eba7ec7b32
25 changed files with 456 additions and 273 deletions
@@ -1,10 +1,7 @@
-use http::header::HeaderName;
 use structopt::StructOpt;
-use tonic::body::BoxBody;
-use tonic::client::GrpcService;
 use tonic::transport::Server;
 use tonic::transport::{Identity, ServerTlsConfig};
-use tonic_interop::{server, MergeTrailers};
+use tonic_interop::server;

 #[derive(StructOpt)]
 struct Opts {
@@ -20,33 +17,7 @@ async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {

    let addr = "127.0.0.1:10000".parse().unwrap();

-    let mut builder = Server::builder().interceptor_fn(|svc, req| {
-        let echo_header = req
-            .headers()
-            .get("x-grpc-test-echo-initial")
-            .map(Clone::clone);
-
-        let echo_trailer = req
-            .headers()
-            .get("x-grpc-test-echo-trailing-bin")
-            .map(Clone::clone)
-            .map(|v| (HeaderName::from_static("x-grpc-test-echo-trailing-bin"), v));
-
-        let call = svc.call(req);
-
-        async move {
-            let mut res = call.await?;
-
-            if let Some(echo_header) = echo_header {
-                res.headers_mut()
-                    .insert("x-grpc-test-echo-initial", echo_header);
-            }
-
-            Ok(res
-                .map(|b| MergeTrailers::new(b, echo_trailer))
-                .map(BoxBody::new))
-        }
-    });
+    let mut builder = Server::builder();

    if matches.use_tls {
        let cert = tokio::fs::read("interop/data/server1.pem").await?;
@@ -60,8 +31,11 @@ async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
    let unimplemented_service =
        server::UnimplementedServiceServer::new(server::UnimplementedService::default());

+    // Wrap this test_service with a service that will echo headers as trailers.
+    let test_service_svc = server::EchoHeadersSvc::new(test_service);
+
    builder
-        .add_service(test_service)
+        .add_service(test_service_svc)
        .add_service(unimplemented_service)
        .serve(addr)
        .await?;
@@ -9,13 +9,7 @@ pub mod pb {
    include!(concat!(env!("OUT_DIR"), "/grpc.testing.rs"));
 }

-use http::header::{HeaderMap, HeaderName, HeaderValue};
-use http_body::Body;
-use std::{
-    default, fmt, iter,
-    pin::Pin,
-    task::{Context, Poll},
-};
+use std::{default, fmt, iter};

 pub fn trace_init() {
    let sub = tracing_subscriber::FmtSubscriber::builder()
@@ -147,41 +141,3 @@ macro_rules! test_assert {
        }
    };
 }
-
-pub struct MergeTrailers<B> {
-    inner: B,
-    trailer: Option<(HeaderName, HeaderValue)>,
-}
-
-impl<B> MergeTrailers<B> {
-    pub fn new(inner: B, trailer: Option<(HeaderName, HeaderValue)>) -> Self {
-        Self { inner, trailer }
-    }
-}
-
-impl<B: Body + Unpin> Body for MergeTrailers<B> {
-    type Data = B::Data;
-    type Error = B::Error;
-
-    fn poll_data(
-        mut self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<Option<Result<Self::Data, Self::Error>>> {
-        Pin::new(&mut self.inner).poll_data(cx)
-    }
-
-    fn poll_trailers(
-        mut self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<Result<Option<HeaderMap>, Self::Error>> {
-        Pin::new(&mut self.inner).poll_trailers(cx).map_ok(|h| {
-            h.map(|mut headers| {
-                if let Some((key, value)) = &self.trailer {
-                    headers.insert(key.clone(), value.clone());
-                }
-
-                headers
-            })
-        })
-    }
-}
@@ -1,9 +1,14 @@
 use crate::pb::{self, *};
 use async_stream::try_stream;
 use futures_util::{stream, StreamExt, TryStreamExt};
+use http::header::{HeaderMap, HeaderName, HeaderValue};
+use http_body::Body;
+use std::future::Future;
 use std::pin::Pin;
+use std::task::{Context, Poll};
 use std::time::Duration;
-use tonic::{Code, Request, Response, Status};
+use tonic::{body::BoxBody, transport::ServiceName, Code, Request, Response, Status};
+use tower::Service;

 pub use pb::test_service_server::TestServiceServer;
 pub use pb::unimplemented_service_server::UnimplementedServiceServer;
@@ -159,3 +164,100 @@ impl pb::unimplemented_service_server::UnimplementedService for UnimplementedSer
        Err(Status::unimplemented(""))
    }
 }
+
+#[derive(Clone, Default)]
+pub struct EchoHeadersSvc<S> {
+    inner: S,
+}
+
+impl<S: ServiceName> ServiceName for EchoHeadersSvc<S> {
+    const NAME: &'static str = S::NAME;
+}
+
+impl<S> EchoHeadersSvc<S> {
+    pub fn new(inner: S) -> Self {
+        Self { inner }
+    }
+}
+
+impl<S> Service<http::Request<hyper::Body>> for EchoHeadersSvc<S>
+where
+    S: Service<http::Request<hyper::Body>, Response = http::Response<BoxBody>> + Send,
+    S::Future: Send + 'static,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = Pin<
+        Box<dyn Future<Output = std::result::Result<Self::Response, Self::Error>> + Send + 'static>,
+    >;
+
+    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
+        Ok(()).into()
+    }
+
+    fn call(&mut self, req: http::Request<hyper::Body>) -> Self::Future {
+        let echo_header = req
+            .headers()
+            .get("x-grpc-test-echo-initial")
+            .map(Clone::clone);
+
+        let echo_trailer = req
+            .headers()
+            .get("x-grpc-test-echo-trailing-bin")
+            .map(Clone::clone)
+            .map(|v| (HeaderName::from_static("x-grpc-test-echo-trailing-bin"), v));
+
+        let call = self.inner.call(req);
+
+        Box::pin(async move {
+            let mut res = call.await?;
+
+            if let Some(echo_header) = echo_header {
+                res.headers_mut()
+                    .insert("x-grpc-test-echo-initial", echo_header);
+            }
+
+            Ok(res
+                .map(|b| MergeTrailers::new(b, echo_trailer))
+                .map(BoxBody::new))
+        })
+    }
+}
+
+pub struct MergeTrailers<B> {
+    inner: B,
+    trailer: Option<(HeaderName, HeaderValue)>,
+}
+
+impl<B> MergeTrailers<B> {
+    pub fn new(inner: B, trailer: Option<(HeaderName, HeaderValue)>) -> Self {
+        Self { inner, trailer }
+    }
+}
+
+impl<B: Body + Unpin> Body for MergeTrailers<B> {
+    type Data = B::Data;
+    type Error = B::Error;
+
+    fn poll_data(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<std::result::Result<Self::Data, Self::Error>>> {
+        Pin::new(&mut self.inner).poll_data(cx)
+    }
+
+    fn poll_trailers(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<std::result::Result<Option<HeaderMap>, Self::Error>> {
+        Pin::new(&mut self.inner).poll_trailers(cx).map_ok(|h| {
+            h.map(|mut headers| {
+                if let Some((key, value)) = &self.trailer {
+                    headers.insert(key.clone(), value.clone());
+                }
+
+                headers
+            })
+        })
+    }
+}