viam_rust_utils/proxy/
grpc_proxy.rs

1use http::uri::{Scheme, Uri};
2use hyper::body::HttpBody;
3use hyper::Request;
4use std::task::{Context, Poll};
5use tower::Service;
6#[derive(Clone, Debug)]
7pub struct GRPCProxy<T> {
8    inner: T,
9    uri: Uri,
10}
11
12impl<T> GRPCProxy<T> {
13    pub fn new(inner: T, uri: Uri) -> Self {
14        GRPCProxy { inner, uri }
15    }
16}
17
18impl<T, ReqBody> Service<Request<ReqBody>> for GRPCProxy<T>
19where
20    T: Service<Request<tonic::body::BoxBody>> + Clone,
21    ReqBody: http_body::Body<Data = hyper::body::Bytes> + Send + 'static,
22    ReqBody::Error: ToString + Send + 'static,
23{
24    type Response = T::Response;
25    type Error = T::Error;
26    type Future = T::Future;
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    fn call(&mut self, request: Request<ReqBody>) -> Self::Future {
32        let (mut h, b) = request.into_parts();
33        let b = b
34            .map_err(|e| tonic::Status::new(tonic::Code::Unknown, e.to_string()))
35            .boxed_unsync();
36        let mut to_uri = self.uri.clone().into_parts();
37        to_uri.path_and_query = h.uri.into_parts().path_and_query;
38        if to_uri.scheme.is_none() {
39            to_uri.scheme = Some(Scheme::HTTPS)
40        }
41        let proxy_uri = Uri::from_parts(to_uri).unwrap();
42        h.uri = proxy_uri;
43        let req = Request::from_parts(h, b);
44        self.inner.call(req)
45    }
46}