app_auth.rs

use crate::{GitHubAppConfig, GitHubError};
use chrono::{DateTime, Duration, Utc};
use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::Mutex;
use tracing::{debug, error, info, instrument};

#[derive(Debug, Serialize)]
struct Claims {
    iat: i64,
    exp: i64,
    iss: String,
}

#[derive(Debug, Deserialize)]
struct AccessTokenResponse {
    token: String,
    expires_at: String,
}

#[derive(Clone)]
struct CachedToken {
    token: String,
    expires_at: DateTime<Utc>,
}

pub struct GitHubTokenProvider {
    config: GitHubAppConfig,
    client: reqwest::Client,
    cached_token: Arc<Mutex<Option<CachedToken>>>,
}

impl GitHubTokenProvider {
    pub fn new(config: GitHubAppConfig) -> Self {
        Self {
            config,
            client: reqwest::Client::new(),
            cached_token: Arc::new(Mutex::new(None)),
        }
    }

    #[instrument(skip(self))]
    fn create_jwt(&self) -> Result<String, GitHubError> {
        debug!("Creating JWT for GitHub App authentication");

        let now = Utc::now();
        let iat = now.timestamp();
        let exp = (now + Duration::minutes(10)).timestamp();

        let claims = Claims {
            iat,
            exp,
            iss: self.config.app_id.to_string(),
        };

        let key = EncodingKey::from_rsa_pem(self.config.private_key_pem.as_bytes())?;
        let token = encode(&Header::new(Algorithm::RS256), &claims, &key)?;

        debug!("JWT created successfully");
        Ok(token)
    }

    #[instrument(skip(self), fields(installation_id = %self.config.installation_id))]
    async fn fetch_installation_token(&self) -> Result<CachedToken, GitHubError> {
        info!("Fetching new installation access token");

        let jwt = self.create_jwt()?;

        let url = format!(
            "https://api.github.com/app/installations/{}/access_tokens",
            self.config.installation_id
        );

        debug!(url = %url, "Requesting installation token");

        let response = self
            .client
            .post(&url)
            .header("Authorization", format!("Bearer {}", jwt))
            .header("Accept", "application/vnd.github+json")
            .header("User-Agent", "github_app")
            .header("X-GitHub-Api-Version", "2022-11-28")
            .send()
            .await?;

        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await?;

            // Tracing will automatically sanitize any tokens in the body
            error!(
                status = %status,
                response_body = %body,
                "Failed to get installation token"
            );

            return Err(GitHubError::Other(format!(
                "Failed to get installation token: {}",
                status
            )));
        }

        let token_response: AccessTokenResponse = response.json().await?;
        let expires_at = DateTime::parse_from_rfc3339(&token_response.expires_at)
            .map_err(|e| GitHubError::Other(format!("Failed to parse expiry time: {}", e)))?
            .with_timezone(&Utc);

        info!(expires_at = %expires_at, "Installation token fetched successfully");

        Ok(CachedToken {
            token: token_response.token,
            expires_at,
        })
    }

    /// Get a valid token, fetching a new one if needed
    ///
    /// This implementation uses double-check locking to minimize contention:
    /// 1. Quick read-only check if token is valid (fast path)
    /// 2. If refresh needed, acquire lock and check again
    /// 3. Only one task fetches new token, others wait and reuse it
    #[instrument(skip(self))]
    pub async fn get_token(&self) -> Result<String, GitHubError> {
        // Fast path: Check if we have a valid token without holding lock during HTTP
        {
            let cached = self.cached_token.lock().await;

            if let Some(token) = &*cached {
                let now = Utc::now();
                let buffer = Duration::minutes(5);

                // Token is still valid
                if token.expires_at - buffer >= now {
                    return Ok(token.token.clone());
                }
            }
            // Lock released here automatically
        }

        // Slow path: Token expired or doesn't exist, need to refresh
        // Acquire lock for the refresh operation
        let mut cached = self.cached_token.lock().await;

        // Double-check: Another task might have refreshed while we waited for lock
        let now = Utc::now();
        let should_refresh = match &*cached {
            None => true,
            Some(token) => {
                let buffer = Duration::minutes(5);
                token.expires_at - buffer < now
            }
        };

        if should_refresh {
            // Lock is held, but we're the only one doing the fetch
            let new_token = self.fetch_installation_token().await?;
            let token_str = new_token.token.clone();
            *cached = Some(new_token);
            Ok(token_str)
        } else {
            // Another task refreshed while we waited
            Ok(cached.as_ref().unwrap().token.clone())
        }
    }
}