volo_http/client/layer/
utils.rs

1use std::borrow::Cow;
2
3use faststr::FastStr;
4use http::{header, uri::Uri};
5use motore::{layer::Layer, service::Service};
6use volo::{client::Apply, context::Context, net::Address};
7
8use crate::{
9    client::Target,
10    context::ClientContext,
11    error::{ClientError, client::Result},
12    request::Request,
13};
14
15/// Set a [`Target`] as the destination forcely.
16///
17/// The layer only sets destination address, including target address, target domain name (for DNS)
18/// and SNI of target host (if using HTTPS).
19///
20/// Note that this layer MUST be set as an outer layer because outer layers runs before service
21/// discover (DNS).
22pub struct TargetLayer {
23    target: Target,
24    service_name: FastStr,
25}
26
27impl TargetLayer {
28    /// Create a [`TargetLayer`] via a [`Target`].
29    ///
30    /// This layer will set the [`Target`] as request destination.
31    pub const fn new(target: Target) -> Self {
32        Self {
33            target,
34            service_name: FastStr::empty(),
35        }
36    }
37
38    /// Create a [`TargetLayer`] via an address.
39    pub fn new_address<A>(addr: A) -> Self
40    where
41        A: Into<Address>,
42    {
43        let addr = addr.into();
44        let target = Target::from(addr);
45        Self {
46            target,
47            service_name: FastStr::empty(),
48        }
49    }
50
51    /// Create a [`TargetLayer`] via a host name.
52    pub fn new_host<S>(host: S) -> Self
53    where
54        S: Into<Cow<'static, str>>,
55    {
56        let target = Target::from_host(host);
57        Self {
58            target,
59            service_name: FastStr::empty(),
60        }
61    }
62
63    /// Create a [`Target`] from a [`Uri`].
64    ///
65    /// ## Panics
66    ///
67    /// This function will panic if the given [`Uri`] cannot be parsed to a [`Target`].
68    pub fn from_uri(uri: &Uri) -> Self {
69        let target = Target::from_uri(uri).expect("invalid uri for building target");
70        Self {
71            target,
72            service_name: FastStr::empty(),
73        }
74    }
75
76    /// Set a service name for current [`TargetLayer`].
77    ///
78    /// When this layer sets [`Target`] as the destination, the service name is also set, and if the
79    /// request uses HTTPS, the service name will be used as the SNI.
80    pub fn with_service_name<S>(mut self, service_name: S) -> Self
81    where
82        S: Into<FastStr>,
83    {
84        self.service_name = service_name.into();
85        self
86    }
87}
88
89impl<S> Layer<S> for TargetLayer {
90    type Service = TargetService<S>;
91
92    fn layer(self, inner: S) -> Self::Service {
93        TargetService {
94            inner,
95            target: self.target,
96            service_name: self.service_name,
97        }
98    }
99}
100
101/// [`Service`] generated by [`TargetLayer`].
102///
103/// For more details, please refer to [`TargetLayer`].
104pub struct TargetService<S> {
105    inner: S,
106    target: Target,
107    service_name: FastStr,
108}
109
110impl<B, S> Service<ClientContext, Request<B>> for TargetService<S>
111where
112    B: Send,
113    S: Service<ClientContext, Request<B>, Error = ClientError> + Send + Sync,
114{
115    type Response = S::Response;
116    type Error = S::Error;
117
118    async fn call(
119        &self,
120        cx: &mut ClientContext,
121        mut req: Request<B>,
122    ) -> Result<Self::Response, Self::Error> {
123        self.target.clone().apply(cx)?;
124        if !self.service_name.is_empty() {
125            cx.rpc_info_mut()
126                .callee_mut()
127                .set_service_name(self.service_name.clone());
128        }
129        // Since `Host` is one of the outermost layers, it cannot know that `Target` has been
130        // modified here, so we need to manually update `Host` here.
131        if !req.headers().contains_key(header::HOST) {
132            if let Some(host) = super::header::gen_host(&self.target) {
133                req.headers_mut().insert(header::HOST, host);
134            }
135        }
136        self.inner.call(cx, req).await
137    }
138}