oauth_device_flows/

token.rs

//! Token management and refresh functionality

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;

/// Manages OAuth tokens including refresh functionality
#[derive(Debug, Clone)]
pub struct TokenManager {
    /// The current token response
    token: TokenResponse,

    /// The OAuth provider
    provider: Provider,

    /// Configuration
    config: DeviceFlowConfig,

    /// HTTP client
    client: Client,
}

impl TokenManager {
    /// Create a new token manager
    pub fn new(token: TokenResponse, provider: Provider, config: DeviceFlowConfig) -> Result<Self> {
        let client = Self::build_client(&config)?;

        Ok(Self {
            token,
            provider,
            config,
            client,
        })
    }

    /// Create a token manager from an existing token (for deserialization)
    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, // Default, should be set properly
            config,
            client,
        }
    }

    /// Get the current access token
    pub fn access_token(&self) -> &str {
        self.token.access_token()
    }

    /// Get the current token response
    pub fn token(&self) -> &TokenResponse {
        &self.token
    }

    /// Check if the token is expired
    pub fn is_expired(&self) -> bool {
        self.token.is_expired()
    }

    /// Check if the token will expire within the given duration
    pub fn expires_within(&self, duration: Duration) -> bool {
        self.token.expires_within(duration)
    }

    /// Get the remaining lifetime of the token
    pub fn remaining_lifetime(&self) -> Option<Duration> {
        self.token.remaining_lifetime()
    }

    /// Refresh the token if a refresh token is available
    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(())
    }

    /// Refresh the token and return the new token without updating the manager
    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, // Keep original scope
        };

        let mut req_builder = self.client.post(token_endpoint).form(&request);

        // Add client secret if required
        if let Some(ref client_secret) = self.config.client_secret {
            req_builder = req_builder.form(&[("client_secret", client_secret.expose_secret())]);
        }

        // Add provider-specific headers
        for (key, value) in self.provider.headers() {
            req_builder = req_builder.header(key, value);
        }

        // Add additional headers
        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?;

        // Update the issued_at timestamp
        token_response.issued_at = OffsetDateTime::now_utc();

        // If no new refresh token was provided, keep the old one
        if token_response.refresh_token.is_none() {
            token_response.refresh_token = self.token.refresh_token.clone();
        }

        Ok(token_response)
    }

    /// Get a valid access token, refreshing if necessary
    pub async fn get_valid_token(&mut self) -> Result<&str> {
        // Check if token is expired or will expire soon (within 5 minutes)
        if self.expires_within(Duration::from_secs(300)) {
            self.refresh().await?;
        }

        Ok(self.access_token())
    }

    /// Create an authorization header value
    pub fn authorization_header(&self) -> String {
        format!("{} {}", self.token.token_type, self.access_token())
    }

    /// Update the provider (useful when deserializing)
    pub fn with_provider(mut self, provider: Provider) -> Self {
        self.provider = provider;
        self
    }

    /// Update the configuration (useful when deserializing)
    pub fn with_config(mut self, config: DeviceFlowConfig) -> Result<Self> {
        self.client = Self::build_client(&config)?;
        self.config = config;
        Ok(self)
    }

    /// Revoke the token (if supported by the provider)
    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(())
    }

    /// Build HTTP client with configuration
    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)
    }
}