tonic/transport/service/
add_origin.rs

1use futures_core::future::BoxFuture;
2use http::{Request, Uri};
3use std::task::{Context, Poll};
4use tower_service::Service;
5
6#[derive(Debug)]
7pub(crate) struct AddOrigin<T> {
8    inner: T,
9    origin: Uri,
10}
11
12impl<T> AddOrigin<T> {
13    pub(crate) fn new(inner: T, origin: Uri) -> Self {
14        Self { inner, origin }
15    }
16}
17
18impl<T, ReqBody> Service<Request<ReqBody>> for AddOrigin<T>
19where
20    T: Service<Request<ReqBody>>,
21    T::Future: Send + 'static,
22    T::Error: Into<crate::Error>,
23{
24    type Response = T::Response;
25    type Error = crate::Error;
26    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
27
28    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
29        self.inner.poll_ready(cx).map_err(Into::into)
30    }
31
32    fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
33        // Split the request into the head and the body.
34        let (mut head, body) = req.into_parts();
35
36        // Split the request URI into parts.
37        let mut uri: http::uri::Parts = head.uri.into();
38        let set_uri = self.origin.clone().into_parts();
39
40        if set_uri.scheme.is_none() || set_uri.authority.is_none() {
41            let err = crate::transport::Error::new_invalid_uri();
42            return Box::pin(async move { Err::<Self::Response, _>(err.into()) });
43        }
44
45        // Update the URI parts, setting hte scheme and authority
46        uri.scheme = Some(set_uri.scheme.expect("expected scheme"));
47        uri.authority = Some(set_uri.authority.expect("expected authority"));
48
49        // Update the the request URI
50        head.uri = http::Uri::from_parts(uri).expect("valid uri");
51
52        let request = Request::from_parts(head, body);
53
54        let fut = self.inner.call(request);
55
56        Box::pin(async move { fut.await.map_err(Into::into) })
57    }
58}