chore(interceptor): Add response body for interceptor

hyperium · Nov 30, 2024 · 60e1d42 · 60e1d42
1 parent 3ad90f6
commit 60e1d42
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 26 deletions.
diff --git a/tonic/src/service/interceptor.rs b/tonic/src/service/interceptor.rs
@@ -116,9 +116,8 @@ impl<S, I, ReqBody, ResBody> Service<http::Request<ReqBody>> for InterceptedServ
 where
     S: Service<http::Request<ReqBody>, Response = http::Response<ResBody>>,
     I: Interceptor,
-    ResBody: Default,
 {
-    type Response = http::Response<ResBody>;
+    type Response = http::Response<ResponseBody<ResBody>>;
     type Error = S::Error;
     type Future = ResponseFuture<S::Future>;
 
@@ -194,21 +193,79 @@ enum Kind<F> {
 impl<F, E, B> Future for ResponseFuture<F>
 where
     F: Future<Output = Result<http::Response<B>, E>>,
-    B: Default,
 {
-    type Output = Result<http::Response<B>, E>;
+    type Output = Result<http::Response<ResponseBody<B>>, E>;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         match self.project().kind.project() {
-            KindProj::Future(future) => future.poll(cx),
+            KindProj::Future(future) => future.poll(cx).map_ok(|res| res.map(ResponseBody::wrap)),
             KindProj::Status(status) => {
-                let response = status.take().unwrap().into_http();
+                let (parts, ()) = status.take().unwrap().into_http::<()>().into_parts();
+                let response = http::Response::from_parts(parts, ResponseBody::<B>::empty());
                 Poll::Ready(Ok(response))
             }
         }
     }
 }
 
+/// Response body for [`InterceptedService`].
+#[pin_project]
+#[derive(Debug)]
+pub struct ResponseBody<B> {
+    #[pin]
+    kind: ResponseBodyKind<B>,
+}
+
+#[pin_project(project = ResponseBodyKindProj)]
+#[derive(Debug)]
+enum ResponseBodyKind<B> {
+    Empty,
+    Wrap(#[pin] B),
+}
+
+impl<B> ResponseBody<B> {
+    fn new(kind: ResponseBodyKind<B>) -> Self {
+        Self { kind }
+    }
+
+    fn empty() -> Self {
+        Self::new(ResponseBodyKind::Empty)
+    }
+
+    fn wrap(body: B) -> Self {
+        Self::new(ResponseBodyKind::Wrap(body))
+    }
+}
+
+impl<B: http_body::Body> http_body::Body for ResponseBody<B> {
+    type Data = B::Data;
+    type Error = B::Error;
+
+    fn poll_frame(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
+        match self.project().kind.project() {
+            ResponseBodyKindProj::Empty => Poll::Ready(None),
+            ResponseBodyKindProj::Wrap(body) => body.poll_frame(cx),
+        }
+    }
+
+    fn size_hint(&self) -> http_body::SizeHint {
+        match &self.kind {
+            ResponseBodyKind::Empty => http_body::SizeHint::with_exact(0),
+            ResponseBodyKind::Wrap(body) => body.size_hint(),
+        }
+    }
+
+    fn is_end_stream(&self) -> bool {
+        match &self.kind {
+            ResponseBodyKind::Empty => true,
+            ResponseBodyKind::Wrap(body) => body.is_end_stream(),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

diff --git a/tonic/src/transport/server/mod.rs b/tonic/src/transport/server/mod.rs
@@ -395,11 +395,8 @@ impl<L> Server<L> {
     /// route around different services.
     pub fn add_service<S>(&mut self, svc: S) -> Router<L>
     where
-        S: Service<Request<Body>, Response = Response<Body>, Error = Infallible>
-            + NamedService
-            + Clone
-            + Send
-            + 'static,
+        S: Service<Request<Body>, Error = Infallible> + NamedService + Clone + Send + 'static,
+        S::Response: axum::response::IntoResponse,
         S::Future: Send + 'static,
         L: Clone,
     {
@@ -416,11 +413,8 @@ impl<L> Server<L> {
     /// As a result, one cannot use this to toggle between two identically named implementations.
     pub fn add_optional_service<S>(&mut self, svc: Option<S>) -> Router<L>
     where
-        S: Service<Request<Body>, Response = Response<Body>, Error = Infallible>
-            + NamedService
-            + Clone
-            + Send
-            + 'static,
+        S: Service<Request<Body>, Error = Infallible> + NamedService + Clone + Send + 'static,
+        S::Response: axum::response::IntoResponse,
         S::Future: Send + 'static,
         L: Clone,
     {
@@ -732,11 +726,8 @@ impl<L> Router<L> {
     /// Add a new service to this router.
     pub fn add_service<S>(mut self, svc: S) -> Self
     where
-        S: Service<Request<Body>, Response = Response<Body>, Error = Infallible>
-            + NamedService
-            + Clone
-            + Send
-            + 'static,
+        S: Service<Request<Body>, Error = Infallible> + NamedService + Clone + Send + 'static,
+        S::Response: axum::response::IntoResponse,
         S::Future: Send + 'static,
     {
         self.routes = self.routes.add_service(svc);
@@ -750,11 +741,8 @@ impl<L> Router<L> {
     /// As a result, one cannot use this to toggle between two identically named implementations.
     pub fn add_optional_service<S>(mut self, svc: Option<S>) -> Self
     where
-        S: Service<Request<Body>, Response = Response<Body>, Error = Infallible>
-            + NamedService
-            + Clone
-            + Send
-            + 'static,
+        S: Service<Request<Body>, Error = Infallible> + NamedService + Clone + Send + 'static,
+        S::Response: axum::response::IntoResponse,
         S::Future: Send + 'static,
     {
         if let Some(svc) = svc {