From 5e5b2e97dca14feeb9bab8568f36ccff72ca608a Mon Sep 17 00:00:00 2001
From: Lucio Franco <luciofranco14@gmail.com>
Date: Wed, 20 Oct 2021 11:02:52 -0400
Subject: [PATCH] fix(transport): AddOrigin panic on invalid uri

---
 tonic/src/transport/service/add_origin.rs | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/tonic/src/transport/service/add_origin.rs b/tonic/src/transport/service/add_origin.rs
index a1d7cd2ca..50d5b6d96 100644
--- a/tonic/src/transport/service/add_origin.rs
+++ b/tonic/src/transport/service/add_origin.rs
@@ -1,3 +1,4 @@
+use futures_core::future::BoxFuture;
 use http::{Request, Uri};
 use std::task::{Context, Poll};
 use tower_service::Service;
@@ -17,13 +18,15 @@ impl<T> AddOrigin<T> {
 impl<T, ReqBody> Service<Request<ReqBody>> for AddOrigin<T>
 where
     T: Service<Request<ReqBody>>,
+    T::Future: Send + 'static,
+    T::Error: Into<crate::Error>,
 {
     type Response = T::Response;
-    type Error = T::Error;
-    type Future = T::Future;
+    type Error = crate::Error;
+    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
 
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        self.inner.poll_ready(cx)
+        self.inner.poll_ready(cx).map_err(Into::into)
     }
 
     fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
@@ -34,6 +37,11 @@ where
         let mut uri: http::uri::Parts = head.uri.into();
         let set_uri = self.origin.clone().into_parts();
 
+        if set_uri.scheme.is_none() || set_uri.authority.is_none() {
+            let err = crate::transport::Error::new_invalid_uri();
+            return Box::pin(async move { Err::<Self::Response, _>(err.into()) });
+        }
+
         // Update the URI parts, setting hte scheme and authority
         uri.scheme = Some(set_uri.scheme.expect("expected scheme"));
         uri.authority = Some(set_uri.authority.expect("expected authority"));
@@ -43,6 +51,8 @@ where
 
         let request = Request::from_parts(head, body);
 
-        self.inner.call(request)
+        let fut = self.inner.call(request);
+
+        Box::pin(async move { fut.await.map_err(Into::into) })
     }
 }