use crate::{
config::DeviceFlowConfig,
error::{DeviceFlowError, Result},
provider::Provider,
types::{RefreshTokenRequest, TokenResponse},
};
use reqwest::Client;
use secrecy::ExposeSecret;
use std::time::Duration;
use time::OffsetDateTime;
use url::Url;
#[derive(Debug, Clone)]
pub struct TokenManager {
token: TokenResponse,
provider: Provider,
config: DeviceFlowConfig,
client: Client,
}
impl TokenManager {
pub fn new(token: TokenResponse, provider: Provider, config: DeviceFlowConfig) -> Result<Self> {
let client = Self::build_client(&config)?;
Ok(Self {
token,
provider,
config,
client,
})
}
pub fn from_token(token: TokenResponse) -> Self {
let config = DeviceFlowConfig::new();
let client = Self::build_client(&config).unwrap_or_default();
Self {
token,
provider: Provider::Microsoft, config,
client,
}
}
pub fn access_token(&self) -> &str {
self.token.access_token()
}
pub fn token(&self) -> &TokenResponse {
&self.token
}
pub fn is_expired(&self) -> bool {
self.token.is_expired()
}
pub fn expires_within(&self, duration: Duration) -> bool {
self.token.expires_within(duration)
}
pub fn remaining_lifetime(&self) -> Option<Duration> {
self.token.remaining_lifetime()
}
pub async fn refresh(&mut self) -> Result<()> {
let refresh_token = self
.token
.refresh_token()
.ok_or_else(|| DeviceFlowError::other("No refresh token available"))?;
let new_token = self.refresh_token(refresh_token).await?;
self.token = new_token;
Ok(())
}
pub async fn refresh_token(&self, refresh_token: &str) -> Result<TokenResponse> {
let token_endpoint = if let Some(ref config) = self.config.generic_provider_config {
config.token_endpoint.clone()
} else {
Url::parse(self.provider.token_endpoint())
.map_err(|e| DeviceFlowError::other(format!("Invalid token endpoint: {e}")))?
};
let request = RefreshTokenRequest {
grant_type: "refresh_token".to_string(),
refresh_token: refresh_token.to_string(),
client_id: self.config.client_id.clone(),
scope: None, };
let mut req_builder = self.client.post(token_endpoint).form(&request);
if let Some(ref client_secret) = self.config.client_secret {
req_builder = req_builder.form(&[("client_secret", client_secret.expose_secret())]);
}
for (key, value) in self.provider.headers() {
req_builder = req_builder.header(key, value);
}
for (key, value) in &self.config.additional_headers {
req_builder = req_builder.header(key, value);
}
let response = req_builder.send().await?;
if !response.status().is_success() {
let error_text = response.text().await?;
return Err(DeviceFlowError::other(format!(
"Token refresh failed: {error_text}"
)));
}
let mut token_response: TokenResponse = response.json().await?;
token_response.issued_at = OffsetDateTime::now_utc();
if token_response.refresh_token.is_none() {
token_response.refresh_token = self.token.refresh_token.clone();
}
Ok(token_response)
}
pub async fn get_valid_token(&mut self) -> Result<&str> {
if self.expires_within(Duration::from_secs(300)) {
self.refresh().await?;
}
Ok(self.access_token())
}
pub fn authorization_header(&self) -> String {
format!("{} {}", self.token.token_type, self.access_token())
}
pub fn with_provider(mut self, provider: Provider) -> Self {
self.provider = provider;
self
}
pub fn with_config(mut self, config: DeviceFlowConfig) -> Result<Self> {
self.client = Self::build_client(&config)?;
self.config = config;
Ok(self)
}
pub async fn revoke(&self) -> Result<()> {
let revoke_endpoint = match self.provider {
Provider::Microsoft => "https://login.microsoftonline.com/common/oauth2/v2.0/logout",
Provider::Google => "https://oauth2.googleapis.com/revoke",
Provider::GitHub => {
return Err(DeviceFlowError::other(
"GitHub does not support token revocation",
))
}
Provider::GitLab => "https://gitlab.com/oauth/revoke",
Provider::Generic => {
return Err(DeviceFlowError::other(
"Revocation not supported for generic provider",
));
}
};
let revoke_url = Url::parse(revoke_endpoint)
.map_err(|e| DeviceFlowError::other(format!("Invalid revoke endpoint: {e}")))?;
let form_data = match self.provider {
Provider::Google => vec![("token", self.access_token())],
Provider::Microsoft => vec![("token", self.access_token())],
Provider::GitLab => vec![
("token", self.access_token()),
("client_id", &self.config.client_id),
],
_ => {
return Err(DeviceFlowError::other(
"Revocation not implemented for this provider",
))
}
};
let response = self.client.post(revoke_url).form(&form_data).send().await?;
if !response.status().is_success() {
let error_text = response.text().await?;
return Err(DeviceFlowError::other(format!(
"Token revocation failed: {error_text}"
)));
}
Ok(())
}
fn build_client(config: &DeviceFlowConfig) -> Result<Client> {
let mut client_builder = Client::builder().timeout(config.request_timeout);
if let Some(ref user_agent) = config.user_agent {
client_builder = client_builder.user_agent(user_agent);
}
client_builder.build().map_err(DeviceFlowError::from)
}
}