Skip to main content

trino_rust_client/auth/
mod.rs

1use std::fmt;
2use std::sync::Arc;
3use std::time::Duration;
4
5mod oauth2;
6pub(crate) use oauth2::run_flow;
7pub use oauth2::{
8    parse_www_authenticate, BrowserRedirectHandler, Challenge, OAuth2State, RedirectHandler,
9};
10
11const DEFAULT_MAX_POLL_ATTEMPTS: usize = 10;
12const DEFAULT_POLL_TIMEOUT: Duration = Duration::from_secs(120);
13
14#[derive(Clone)]
15#[non_exhaustive]
16pub enum Auth {
17    Basic(String, Option<String>),
18    Jwt(String),
19    OAuth2(Arc<OAuth2State>),
20}
21
22impl Auth {
23    pub fn new_basic(username: impl ToString, password: Option<impl ToString>) -> Auth {
24        Auth::Basic(username.to_string(), password.map(|p| p.to_string()))
25    }
26
27    pub fn new_jwt(token: impl ToString) -> Auth {
28        Auth::Jwt(token.to_string())
29    }
30
31    /// Interactive OAuth2 using the default browser handler.
32    pub fn new_oauth2() -> Auth {
33        Auth::new_oauth2_with_handler(Arc::new(BrowserRedirectHandler))
34    }
35
36    /// Interactive OAuth2 with a caller-supplied redirect handler.
37    pub fn new_oauth2_with_handler(handler: Arc<dyn RedirectHandler>) -> Auth {
38        Auth::OAuth2(Arc::new(OAuth2State::new(
39            handler,
40            DEFAULT_MAX_POLL_ATTEMPTS,
41            DEFAULT_POLL_TIMEOUT,
42        )))
43    }
44
45    /// Override the token-server poll settings. No-op for non-OAuth2 auth.
46    pub fn with_poll(self, max_attempts: usize, timeout: Duration) -> Auth {
47        match self {
48            Auth::OAuth2(state) => Auth::OAuth2(Arc::new(OAuth2State::new(
49                Arc::clone(&state.handler),
50                max_attempts,
51                timeout,
52            ))),
53            other => other,
54        }
55    }
56}
57
58impl fmt::Debug for Auth {
59    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
60        match self {
61            Auth::Basic(name, _) => f
62                .debug_struct("BasicAuth")
63                .field("username", name)
64                .field("password", &"******")
65                .finish(),
66
67            Auth::Jwt(_) => f.debug_struct("JwtAuth").field("token", &"******").finish(),
68
69            Auth::OAuth2(_) => f
70                .debug_struct("OAuth2Auth")
71                .field("token", &"******")
72                .finish(),
73        }
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80
81    #[test]
82    fn oauth2_debug_redacts_token() {
83        let auth = Auth::new_oauth2();
84        if let Auth::OAuth2(state) = &auth {
85            *state.token.write().unwrap() = Some("super-secret".to_string());
86        }
87        let dbg = format!("{auth:?}");
88        assert!(!dbg.contains("super-secret"), "token leaked: {dbg}");
89        assert!(dbg.contains("OAuth2Auth"));
90    }
91}