From 62f0285ff0eb478488e737e352e4057e20296f3a Mon Sep 17 00:00:00 2001
From: Lucio Franco <luciofranco14@gmail.com>
Date: Sat, 21 Sep 2019 10:23:31 -0600
Subject: [PATCH] implemented server cusotm metadata

---
 tonic-build/src/service.rs      |  2 +-
 tonic-interop/Cargo.toml        |  1 +
 tonic-interop/src/bin/server.rs | 24 +++++++++++++----
 tonic-interop/src/lib.rs        | 46 ++++++++++++++++++++++++++++++++-
 tonic-interop/test.sh           |  2 +-
 5 files changed, 67 insertions(+), 8 deletions(-)

diff --git a/tonic-build/src/service.rs b/tonic-build/src/service.rs
index d6889ec..96c39b0 100644
--- a/tonic-build/src/service.rs
+++ b/tonic-build/src/service.rs
@@ -67,7 +67,7 @@ pub(crate) fn generate(service: &Service, proto_path: &str) -> TokenStream {
                     #methods
 
                     // TODO: implement grpc unimplemented for server
-                    _ => Err(tonic::Status::unimplemented("")),
+                    _ => unimplemented!(),
                 }
             }
         }
diff --git a/tonic-interop/Cargo.toml b/tonic-interop/Cargo.toml
index 9c4c86f..a215408 100644
--- a/tonic-interop/Cargo.toml
+++ b/tonic-interop/Cargo.toml
@@ -23,6 +23,7 @@ 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"
+http-body = "0.2.0-alpha.1"
 
 console = "0.7"
 structopt = "0.2"
diff --git a/tonic-interop/src/bin/server.rs b/tonic-interop/src/bin/server.rs
index 83f228c..7c8a748 100644
--- a/tonic-interop/src/bin/server.rs
+++ b/tonic-interop/src/bin/server.rs
@@ -1,9 +1,11 @@
 use structopt::StructOpt;
 use tonic::Server;
-use tonic_interop::server;
+use tonic_interop::{server, MergeTrailers};
 // TODO: move GrpcService out of client since it can be used for the
 // server too.
 use tonic::client::GrpcService;
+use tonic::body::BoxBody;
+use http::header::HeaderName;
 
 #[derive(StructOpt)]
 struct Opts {
@@ -30,12 +32,24 @@ async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
     }
 
     builder.interceptor_fn(|svc, req| {
-        println!("INBOUND REQUEST={:?}", 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 res = call.await?;
-            println!("OUTBOUND RESPONSE={:?}", res);
-            Ok(res)
+            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))
         }
     });
 
diff --git a/tonic-interop/src/lib.rs b/tonic-interop/src/lib.rs
index 6239a1b..d744136 100644
--- a/tonic-interop/src/lib.rs
+++ b/tonic-interop/src/lib.rs
@@ -9,7 +9,13 @@ pub mod pb {
     include!(concat!(env!("OUT_DIR"), "/grpc.testing.rs"));
 }
 
-use std::{default, fmt, iter};
+use http::header::{HeaderMap, HeaderName, HeaderValue};
+use http_body::Body;
+use std::{
+    default, fmt, iter,
+    pin::Pin,
+    task::{Context, Poll},
+};
 
 pub fn trace_init() {
     let sub = tracing_subscriber::FmtSubscriber::builder()
@@ -141,3 +147,41 @@ 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
+            })
+        })
+    }
+}
diff --git a/tonic-interop/test.sh b/tonic-interop/test.sh
index 8c58cdc..598c020 100755
--- a/tonic-interop/test.sh
+++ b/tonic-interop/test.sh
@@ -49,6 +49,6 @@ sleep 1
 
 ./target/debug/client \
 --test_case=empty_unary,large_unary,client_streaming,server_streaming,ping_pong,\
-empty_stream,status_code_and_message,special_status_message,unimplemented_method $ARG
+empty_stream,status_code_and_message,special_status_message,unimplemented_method,custom_metadata $ARG
 # ,unimplemented_service,custom_metadata