1use std::fmt;
11use std::future::Future;
12use std::pin::Pin;
13use std::sync::Arc;
14use std::time::Duration;
15
16use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity};
17
18use crate::retry::RetryConfig;
19use crate::{Error, Result};
20
21pub trait TokenSource: Send + Sync + fmt::Debug {
27 fn fetch_bearer(&self) -> Pin<Box<dyn Future<Output = Result<Option<String>>> + Send + '_>>;
30}
31
32#[derive(Clone)]
38#[non_exhaustive]
39pub enum Auth {
40 None,
42 Static(String),
44 Dynamic(Arc<dyn TokenSource>),
46}
47
48impl Auth {
49 pub async fn bearer(&self) -> Result<Option<String>> {
54 match self {
55 Auth::None => Ok(None),
56 Auth::Static(token) => Ok(Some(token.clone())),
57 Auth::Dynamic(source) => source.fetch_bearer().await,
58 }
59 }
60}
61
62impl fmt::Debug for Auth {
63 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64 match self {
65 Auth::None => f.write_str("None"),
66 Auth::Static(_) => f.write_str("Static(<redacted>)"),
67 Auth::Dynamic(source) => write!(f, "Dynamic({source:?})"),
68 }
69 }
70}
71
72#[derive(Clone, Debug, Default)]
79#[non_exhaustive]
80pub struct TlsConfig {
81 pub ca_certificate_pem: Option<Vec<u8>>,
84 pub domain_name: Option<String>,
87 pub client_identity_pem: Option<(Vec<u8>, Vec<u8>)>,
89}
90
91impl TlsConfig {
92 #[must_use]
94 pub fn new() -> Self {
95 Self::default()
96 }
97
98 #[must_use]
100 pub fn with_ca_certificate(mut self, ca_pem: impl Into<Vec<u8>>) -> Self {
101 self.ca_certificate_pem = Some(ca_pem.into());
102 self
103 }
104
105 #[must_use]
107 pub fn with_domain_name(mut self, domain: impl Into<String>) -> Self {
108 self.domain_name = Some(domain.into());
109 self
110 }
111
112 #[must_use]
114 pub fn with_client_identity(
115 mut self,
116 certificate_pem: impl Into<Vec<u8>>,
117 private_key_pem: impl Into<Vec<u8>>,
118 ) -> Self {
119 self.client_identity_pem = Some((certificate_pem.into(), private_key_pem.into()));
120 self
121 }
122}
123
124fn build_tls(tls: Option<&TlsConfig>) -> ClientTlsConfig {
127 let mut config = ClientTlsConfig::new();
128 match tls.and_then(|t| t.ca_certificate_pem.as_ref()) {
129 Some(ca) => config = config.ca_certificate(Certificate::from_pem(ca.clone())),
130 None => config = config.with_native_roots(),
131 }
132 if let Some(domain) = tls.and_then(|t| t.domain_name.as_ref()) {
133 config = config.domain_name(domain.clone());
134 }
135 if let Some((cert, key)) = tls.and_then(|t| t.client_identity_pem.as_ref()) {
136 config = config.identity(Identity::from_pem(cert.clone(), key.clone()));
137 }
138 config
139}
140
141#[derive(Clone, Debug)]
144pub struct Config {
145 endpoint: String,
146 auth: Auth,
147 retry: Option<RetryConfig>,
148 tls: Option<TlsConfig>,
149 timeout: Option<Duration>,
150}
151
152impl Config {
153 pub fn new(endpoint: impl Into<String>) -> Self {
156 Self {
157 endpoint: endpoint.into(),
158 auth: Auth::None,
159 retry: None,
160 tls: None,
161 timeout: None,
162 }
163 }
164
165 #[must_use]
167 pub fn with_token(mut self, token: impl Into<String>) -> Self {
168 self.auth = Auth::Static(token.into());
169 self
170 }
171
172 #[must_use]
175 pub fn with_oidc<T: TokenSource + 'static>(mut self, provider: T) -> Self {
176 self.auth = Auth::Dynamic(Arc::new(provider));
177 self
178 }
179
180 #[must_use]
182 pub fn with_retry(mut self, retry: RetryConfig) -> Self {
183 self.retry = Some(retry);
184 self
185 }
186
187 #[must_use]
189 pub fn with_tls(mut self, tls: TlsConfig) -> Self {
190 self.tls = Some(tls);
191 self
192 }
193
194 #[must_use]
197 pub fn with_timeout(mut self, timeout: Duration) -> Self {
198 self.timeout = Some(timeout);
199 self
200 }
201
202 #[must_use]
204 pub fn endpoint(&self) -> &str {
205 &self.endpoint
206 }
207
208 #[must_use]
210 pub fn auth(&self) -> &Auth {
211 &self.auth
212 }
213
214 #[must_use]
216 pub fn retry(&self) -> Option<&RetryConfig> {
217 self.retry.as_ref()
218 }
219
220 pub fn connect_channel(&self) -> Result<Channel> {
231 let (uri, want_tls) = resolve_endpoint(&self.endpoint, self.tls.is_some());
232 let mut endpoint = Endpoint::from_shared(uri.clone())
233 .map_err(|e| Error::InvalidRequest(format!("invalid endpoint uri {uri:?}: {e}")))?
234 .timeout(self.timeout.unwrap_or(Duration::from_secs(30)))
235 .connect_timeout(Duration::from_secs(10))
236 .http2_keep_alive_interval(Duration::from_secs(30))
237 .keep_alive_timeout(Duration::from_secs(20))
238 .keep_alive_while_idle(true)
239 .tcp_keepalive(Some(Duration::from_secs(60)))
240 .tcp_nodelay(true);
241
242 if want_tls {
243 endpoint = endpoint
244 .tls_config(build_tls(self.tls.as_ref()))
245 .map_err(|e| Error::InvalidRequest(format!("invalid TLS config: {e}")))?;
246 }
247
248 Ok(endpoint.connect_lazy())
249 }
250}
251
252fn resolve_endpoint(endpoint: &str, tls_configured: bool) -> (String, bool) {
262 let is_https = endpoint
263 .get(..8)
264 .is_some_and(|s| s.eq_ignore_ascii_case("https://"));
265 let is_http = endpoint
266 .get(..7)
267 .is_some_and(|s| s.eq_ignore_ascii_case("http://"));
268 let want_tls = tls_configured || is_https;
269 if want_tls && is_http {
270 (format!("https://{}", &endpoint[7..]), true)
271 } else {
272 (endpoint.to_string(), want_tls)
273 }
274}
275
276#[cfg(test)]
277#[allow(clippy::bool_assert_comparison)]
278mod tests {
279 use super::resolve_endpoint;
280
281 #[test]
282 fn with_tls_on_http_endpoint_is_upgraded_to_https() {
283 let (uri, tls) = resolve_endpoint("http://host:5001", true);
286 assert_eq!(uri, "https://host:5001");
287 assert_eq!(tls, true);
288 }
289
290 #[test]
291 fn https_scheme_detection_is_case_insensitive() {
292 let (uri, tls) = resolve_endpoint("HTTPS://host:443", false);
295 assert_eq!(uri, "HTTPS://host:443");
296 assert_eq!(tls, true);
297 }
298
299 #[test]
300 fn plain_http_without_tls_stays_plaintext() {
301 let (uri, tls) = resolve_endpoint("http://host:3901", false);
302 assert_eq!(uri, "http://host:3901");
303 assert_eq!(tls, false);
304 }
305
306 #[test]
307 fn https_without_explicit_tls_wants_tls() {
308 let (uri, tls) = resolve_endpoint("https://host:443", false);
309 assert_eq!(uri, "https://host:443");
310 assert_eq!(tls, true);
311 }
312
313 #[test]
314 fn uppercase_http_with_tls_is_upgraded() {
315 let (uri, tls) = resolve_endpoint("HTTP://host:5001", true);
316 assert_eq!(uri, "https://host:5001");
317 assert_eq!(tls, true);
318 }
319}