Skip to main content

oauth_token_service/
lib.rs

1use base64::engine::general_purpose::URL_SAFE;
2use base64::Engine;
3use log::error;
4use oauth2::basic::{BasicClient, BasicTokenResponse};
5use oauth2::AuthType::RequestBody;
6use oauth2::{
7    reqwest, AccessToken, AuthUrl, ClientId, ClientSecret, HttpClientError, RequestTokenError,
8    Scope, TokenResponse, TokenUrl,
9};
10use std::error::Error;
11use std::fmt::Debug;
12use std::sync::Arc;
13use std::time::{Duration, SystemTime};
14use tokio::sync::Mutex;
15
16#[derive(Debug, thiserror::Error)]
17pub enum TokenServiceError {
18    #[error(transparent)]
19    TokenError(#[from] Box<dyn Error>),
20    #[error(transparent)]
21    NetworkError(#[from] HttpClientError<reqwest::Error>),
22}
23
24/// Information about a token, including the token itself and when it expires
25#[derive(Clone, Debug)]
26pub struct TokenInfo {
27    pub access_token: AccessToken,
28    pub expires_at: SystemTime,
29}
30
31#[derive(Clone, Debug)]
32pub struct TokenServiceConfig {
33    pub identity_service_base_url: String,
34    pub username: String,
35    pub token: String,
36    pub client_id: String,
37}
38
39/// A connector to the identity service that auto-renews its token when expired
40#[derive(Debug, Clone)]
41pub struct TokenService {
42    config: TokenServiceConfig,
43    token_info: Arc<Mutex<Option<TokenInfo>>>,
44}
45
46impl TokenService {
47    pub fn new(config: TokenServiceConfig) -> Self {
48        Self {
49            config,
50            token_info: Arc::new(Mutex::new(None)),
51        }
52    }
53
54    async fn initialize_service(&self) -> Result<TokenInfo, TokenServiceError> {
55        let token = self.perform_login().await?;
56
57        let expires_at = SystemTime::now()
58            + Duration::from_secs(
59                token
60                    .expires_in()
61                    .ok_or_else(|| {
62                        TokenServiceError::TokenError("Token has no duration".to_string().into())
63                    })?
64                    .as_secs(),
65            );
66
67        Ok(TokenInfo {
68            access_token: token.access_token().clone(),
69            expires_at,
70        })
71    }
72
73    async fn perform_login(&self) -> Result<BasicTokenResponse, TokenServiceError> {
74        let client_secret =
75            URL_SAFE.encode(format!("{}:{}", self.config.username, self.config.token));
76        let client = BasicClient::new(ClientId::new(self.config.client_id.clone()))
77            .set_auth_type(RequestBody)
78            .set_client_secret(ClientSecret::new(client_secret))
79            .set_auth_uri(
80                AuthUrl::new(format!(
81                    "{}/authorize/",
82                    self.config.identity_service_base_url
83                ))
84                .expect("Auth URL should be valid"),
85            )
86            .set_token_uri(
87                TokenUrl::new(format!("{}/token/", self.config.identity_service_base_url))
88                    .expect("Token URL should be valid"),
89            );
90
91        let http_client = reqwest::ClientBuilder::new()
92            // Following redirects opens the client up to SSRF vulnerabilities.
93            .redirect(reqwest::redirect::Policy::none())
94            .build()
95            .expect("Client should build");
96
97        let result = client
98            .exchange_client_credentials()
99            .add_scope(Scope::new("profile".to_string()))
100            .request_async(&http_client)
101            .await;
102
103        result.map_err(|e| match e {
104            RequestTokenError::Request(e) => TokenServiceError::NetworkError(e),
105            RequestTokenError::ServerResponse(e) => {
106                TokenServiceError::TokenError(e.to_string().into())
107            }
108            _ => {
109                error!("Unexpected error: {:?}", e);
110                TokenServiceError::TokenError("Unexpected error".to_string().into())
111            }
112        })
113    }
114
115    pub async fn get_token(&self) -> Result<AccessToken, TokenServiceError> {
116        let mut token_info = self.token_info.lock().await;
117
118        if token_info.is_none() || token_info.as_ref().unwrap().expires_at < SystemTime::now() {
119            let new_token_info = self.initialize_service().await?;
120            *token_info = Some(new_token_info);
121        }
122
123        Ok(token_info.as_ref().unwrap().access_token.clone())
124    }
125}