sui_rpc/client/
interceptors.rs

1/// Interceptor used to add additional headers to a Request
2#[derive(Clone, Debug, Default)]
3pub struct HeadersInterceptor {
4    headers: tonic::metadata::MetadataMap,
5}
6
7impl HeadersInterceptor {
8    /// Create a new, empty `HeadersInterceptor`.
9    pub fn new() -> Self {
10        Self::default()
11    }
12
13    /// Return reference to the internal `MetadataMap`.
14    pub fn headers(&self) -> &tonic::metadata::MetadataMap {
15        &self.headers
16    }
17
18    /// Get mutable access to the internal `MetadataMap` for modification.
19    pub fn headers_mut(&mut self) -> &mut tonic::metadata::MetadataMap {
20        &mut self.headers
21    }
22
23    /// Enable HTTP basic authentication with a username and optional password.
24    pub fn basic_auth<U, P>(&mut self, username: U, password: Option<P>)
25    where
26        U: std::fmt::Display,
27        P: std::fmt::Display,
28    {
29        use base64::prelude::BASE64_STANDARD;
30        use base64::write::EncoderWriter;
31        use std::io::Write;
32
33        let mut buf = b"Basic ".to_vec();
34        {
35            let mut encoder = EncoderWriter::new(&mut buf, &BASE64_STANDARD);
36            let _ = write!(encoder, "{username}:");
37            if let Some(password) = password {
38                let _ = write!(encoder, "{password}");
39            }
40        }
41        let mut header = tonic::metadata::MetadataValue::try_from(buf)
42            .expect("base64 is always valid HeaderValue");
43        header.set_sensitive(true);
44
45        self.headers
46            .insert(http::header::AUTHORIZATION.as_str(), header);
47    }
48
49    /// Enable HTTP bearer authentication.
50    pub fn bearer_auth<T>(&mut self, token: T)
51    where
52        T: std::fmt::Display,
53    {
54        let header_value = format!("Bearer {token}");
55        let mut header = tonic::metadata::MetadataValue::try_from(header_value)
56            .expect("token is always valid HeaderValue");
57        header.set_sensitive(true);
58
59        self.headers
60            .insert(http::header::AUTHORIZATION.as_str(), header);
61    }
62}
63
64impl tonic::service::Interceptor for &HeadersInterceptor {
65    fn call(
66        &mut self,
67        mut request: tonic::Request<()>,
68    ) -> std::result::Result<tonic::Request<()>, tonic::Status> {
69        if !self.headers.is_empty() {
70            request
71                .metadata_mut()
72                .as_mut()
73                .extend(self.headers.clone().into_headers());
74        }
75        Ok(request)
76    }
77}
78
79impl tonic::service::Interceptor for HeadersInterceptor {
80    fn call(
81        &mut self,
82        request: tonic::Request<()>,
83    ) -> std::result::Result<tonic::Request<()>, tonic::Status> {
84        (&*self).call(request)
85    }
86}