use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity};
use crate::retry::RetryConfig;
use crate::{Error, Result};
pub trait TokenSource: Send + Sync + fmt::Debug {
fn fetch_bearer(&self) -> Pin<Box<dyn Future<Output = Result<Option<String>>> + Send + '_>>;
}
#[derive(Clone)]
#[non_exhaustive]
pub enum Auth {
None,
Static(String),
Dynamic(Arc<dyn TokenSource>),
}
impl Auth {
pub async fn bearer(&self) -> Result<Option<String>> {
match self {
Auth::None => Ok(None),
Auth::Static(token) => Ok(Some(token.clone())),
Auth::Dynamic(source) => source.fetch_bearer().await,
}
}
}
impl fmt::Debug for Auth {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Auth::None => f.write_str("None"),
Auth::Static(_) => f.write_str("Static(<redacted>)"),
Auth::Dynamic(source) => write!(f, "Dynamic({source:?})"),
}
}
}
#[derive(Clone, Debug, Default)]
#[non_exhaustive]
pub struct TlsConfig {
pub ca_certificate_pem: Option<Vec<u8>>,
pub domain_name: Option<String>,
pub client_identity_pem: Option<(Vec<u8>, Vec<u8>)>,
}
impl TlsConfig {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_ca_certificate(mut self, ca_pem: impl Into<Vec<u8>>) -> Self {
self.ca_certificate_pem = Some(ca_pem.into());
self
}
#[must_use]
pub fn with_domain_name(mut self, domain: impl Into<String>) -> Self {
self.domain_name = Some(domain.into());
self
}
#[must_use]
pub fn with_client_identity(
mut self,
certificate_pem: impl Into<Vec<u8>>,
private_key_pem: impl Into<Vec<u8>>,
) -> Self {
self.client_identity_pem = Some((certificate_pem.into(), private_key_pem.into()));
self
}
}
fn build_tls(tls: Option<&TlsConfig>) -> ClientTlsConfig {
let mut config = ClientTlsConfig::new();
match tls.and_then(|t| t.ca_certificate_pem.as_ref()) {
Some(ca) => config = config.ca_certificate(Certificate::from_pem(ca.clone())),
None => config = config.with_native_roots(),
}
if let Some(domain) = tls.and_then(|t| t.domain_name.as_ref()) {
config = config.domain_name(domain.clone());
}
if let Some((cert, key)) = tls.and_then(|t| t.client_identity_pem.as_ref()) {
config = config.identity(Identity::from_pem(cert.clone(), key.clone()));
}
config
}
#[derive(Clone, Debug)]
pub struct Config {
endpoint: String,
auth: Auth,
retry: Option<RetryConfig>,
tls: Option<TlsConfig>,
timeout: Option<Duration>,
}
impl Config {
pub fn new(endpoint: impl Into<String>) -> Self {
Self {
endpoint: endpoint.into(),
auth: Auth::None,
retry: None,
tls: None,
timeout: None,
}
}
#[must_use]
pub fn with_token(mut self, token: impl Into<String>) -> Self {
self.auth = Auth::Static(token.into());
self
}
#[must_use]
pub fn with_oidc<T: TokenSource + 'static>(mut self, provider: T) -> Self {
self.auth = Auth::Dynamic(Arc::new(provider));
self
}
#[must_use]
pub fn with_retry(mut self, retry: RetryConfig) -> Self {
self.retry = Some(retry);
self
}
#[must_use]
pub fn with_tls(mut self, tls: TlsConfig) -> Self {
self.tls = Some(tls);
self
}
#[must_use]
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.timeout = Some(timeout);
self
}
#[must_use]
pub fn endpoint(&self) -> &str {
&self.endpoint
}
#[must_use]
pub fn auth(&self) -> &Auth {
&self.auth
}
#[must_use]
pub fn retry(&self) -> Option<&RetryConfig> {
self.retry.as_ref()
}
pub fn connect_channel(&self) -> Result<Channel> {
let mut endpoint = Endpoint::from_shared(self.endpoint.clone())
.map_err(|e| {
Error::InvalidRequest(format!("invalid endpoint uri {:?}: {e}", self.endpoint))
})?
.timeout(self.timeout.unwrap_or(Duration::from_secs(30)))
.connect_timeout(Duration::from_secs(10))
.http2_keep_alive_interval(Duration::from_secs(30))
.keep_alive_timeout(Duration::from_secs(20))
.keep_alive_while_idle(true)
.tcp_keepalive(Some(Duration::from_secs(60)))
.tcp_nodelay(true);
if self.tls.is_some() || self.endpoint.starts_with("https") {
endpoint = endpoint
.tls_config(build_tls(self.tls.as_ref()))
.map_err(|e| Error::InvalidRequest(format!("invalid TLS config: {e}")))?;
}
Ok(endpoint.connect_lazy())
}
}