systemprompt-oauth 0.13.0

//! JWT validation provider implementations.

use systemprompt_models::auth::{AuthenticatedUser, JwtAudience, Permission};
use systemprompt_traits::{
    AgentJwtClaims, GenerateTokenParams, JwtProviderError, JwtResult, JwtValidationProvider,
};
use uuid::Uuid;

use super::generation::{JwtConfig, JwtSigningParams, generate_jwt, generate_secure_token};
use super::validation::jwt::validate_jwt_token;

#[derive(Debug)]
pub struct JwtValidationProviderImpl {
    issuer: String,
    audiences: Vec<JwtAudience>,
}

impl JwtValidationProviderImpl {
    #[must_use]
    pub const fn new(issuer: String, audiences: Vec<JwtAudience>) -> Self {
        Self { issuer, audiences }
    }

    pub fn from_config() -> JwtResult<Self> {
        let config = systemprompt_models::Config::get()
            .map_err(|e| JwtProviderError::ConfigurationError(e.to_string()))?;

        Ok(Self {
            issuer: config.jwt_issuer.clone(),
            audiences: config.jwt_audiences.clone(),
        })
    }
}

impl JwtValidationProvider for JwtValidationProviderImpl {
    fn validate_token(&self, token: &str) -> JwtResult<AgentJwtClaims> {
        let claims = validate_jwt_token(token, &self.issuer, &self.audiences).map_err(|e| {
            if e.to_string().contains("expired") {
                JwtProviderError::TokenExpired
            } else {
                JwtProviderError::InvalidToken
            }
        })?;

        let is_admin = claims.is_admin();
        Ok(AgentJwtClaims {
            subject: claims.sub,
            username: claims.username,
            user_type: claims.user_type.to_string(),
            audiences: claims.aud.iter().map(ToString::to_string).collect(),
            permissions: claims.scope.iter().map(ToString::to_string).collect(),
            is_admin,
            expires_at: claims.exp,
            issued_at: claims.iat,
        })
    }

    fn generate_token(&self, params: GenerateTokenParams) -> JwtResult<String> {
        let user_id = Uuid::parse_str(params.user_id.as_str()).map_err(|e| {
            JwtProviderError::Internal(format!(
                "user_id {:?} is not a valid UUID: {e}",
                params.user_id.as_str()
            ))
        })?;

        let user = AuthenticatedUser {
            id: user_id,
            username: params.username.clone(),
            email: params.username.clone(),
            roles: vec![],
            permissions: vec![],
            attributes: std::collections::BTreeMap::new(),
        };

        let permissions: Vec<Permission> = params
            .permissions
            .iter()
            .filter_map(|p| p.parse().ok())
            .collect();

        let audiences: Vec<JwtAudience> = params
            .audiences
            .iter()
            .filter_map(|a| a.parse().ok())
            .collect();

        let config = JwtConfig {
            permissions,
            audience: if audiences.is_empty() {
                JwtAudience::standard()
            } else {
                audiences
            },
            expires_in_hours: params.expires_in_hours.map(i64::from),
            resource: None,
            plugin_id: None,
        };

        let jti = generate_secure_token("jwt");
        let signing = JwtSigningParams {
            issuer: &self.issuer,
        };

        generate_jwt(&user, config, jti, &params.session_id, &signing)
            .map_err(|e| JwtProviderError::Internal(e.to_string()))
    }

    fn generate_secure_token(&self, prefix: &str) -> String {
        generate_secure_token(prefix)
    }
}