mod jwt;
mod oauth;
mod permissions;
mod session;
mod security;
mod revocation;
pub use jwt::{RiJWTManager, RiJWTClaims, RiJWTValidationOptions};
pub use oauth::{RiOAuthManager, RiOAuthToken, RiOAuthUserInfo, RiOAuthProvider};
pub use permissions::{RiPermissionManager, RiPermission, RiRole};
pub use session::{RiSessionManager, RiSession};
pub use security::RiSecurityManager;
pub use revocation::{RiJWTRevocationList, RiRevokedTokenInfo};
use crate::core::{RiResult, RiError, RiServiceContext};
use rand::RngCore;
use serde::Deserialize;
use serde::Serialize;
use std::collections::HashMap;
use std::env;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::RwLock;
#[cfg(feature = "pyo3")]
use tokio::runtime::Handle;
const DEFAULT_JWT_SECRET_ENV: &str = "Ri_JWT_SECRET";
const FALLBACK_SECRET_LENGTH: usize = 64;
#[derive(Debug, Clone)]
struct LoginAttempt {
count: u32,
first_attempt: u64,
locked_until: Option<u64>,
}
#[derive(Debug, Clone)]
pub struct RiRateLimiter {
attempts: Arc<RwLock<HashMap<String, LoginAttempt>>>,
max_attempts: u32,
lockout_secs: u64,
window_secs: u64,
}
impl RiRateLimiter {
pub fn new(max_attempts: u32, lockout_secs: u64, window_secs: u64) -> Self {
Self {
attempts: Arc::new(RwLock::new(HashMap::new())),
max_attempts,
lockout_secs,
window_secs,
}
}
pub async fn check_and_record(&self, identifier: &str) -> Result<(), u64> {
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
let mut attempts = self.attempts.write().await;
if let Some(attempt) = attempts.get_mut(identifier) {
if let Some(locked_until) = attempt.locked_until {
if now < locked_until {
return Err(locked_until - now);
}
attempts.remove(identifier);
return Ok(());
}
if now - attempt.first_attempt > self.window_secs {
attempts.remove(identifier);
return Ok(());
}
attempt.count += 1;
if attempt.count >= self.max_attempts {
attempt.locked_until = Some(now + self.lockout_secs);
return Err(self.lockout_secs);
}
} else {
attempts.insert(identifier.to_string(), LoginAttempt {
count: 1,
first_attempt: now,
locked_until: None,
});
}
Ok(())
}
pub async fn reset(&self, identifier: &str) {
self.attempts.write().await.remove(identifier);
}
}
#[derive(Debug, Clone, Serialize)]
pub struct RiAuditEvent {
pub timestamp: String,
pub event_type: String,
pub user_identifier: Option<String>,
pub ip_address: Option<String>,
pub action: String,
pub resource: Option<String>,
pub success: bool,
pub details: Option<String>,
}
impl RiAuditEvent {
pub fn new(event_type: &str, action: &str) -> Self {
Self {
timestamp: chrono::Utc::now().to_rfc3339(),
event_type: event_type.to_string(),
user_identifier: None,
ip_address: None,
action: action.to_string(),
resource: None,
success: true,
details: None,
}
}
pub fn with_user(mut self, user: &str) -> Self {
self.user_identifier = Some(user.to_string());
self
}
pub fn with_ip(mut self, ip: &str) -> Self {
self.ip_address = Some(ip.to_string());
self
}
pub fn with_resource(mut self, resource: &str) -> Self {
self.resource = Some(resource.to_string());
self
}
pub fn with_details(mut self, details: &str) -> Self {
self.details = Some(details.to_string());
self
}
pub fn with_success(mut self, success: bool) -> Self {
self.success = success;
self
}
}
#[derive(Debug, Clone)]
pub struct RiAuditLogger {
events: Arc<RwLock<Vec<RiAuditEvent>>>,
max_events: usize,
}
impl RiAuditLogger {
pub fn new(max_events: usize) -> Self {
Self {
events: Arc::new(RwLock::new(Vec::with_capacity(max_events))),
max_events,
}
}
pub async fn log(&self, event: RiAuditEvent) {
let mut events = self.events.write().await;
if events.len() >= self.max_events {
events.remove(0);
}
events.push(event);
}
pub async fn get_events(&self) -> Vec<RiAuditEvent> {
self.events.read().await.clone()
}
}
fn load_jwt_secret_from_env() -> String {
env::var(DEFAULT_JWT_SECRET_ENV).unwrap_or_else(|_| {
let mut secret = vec![0u8; FALLBACK_SECRET_LENGTH];
rand::thread_rng().fill_bytes(&mut secret);
hex::encode(secret)
})
}
fn load_oauth_env_var(provider_name: &str, suffix: &str) -> Result<String, RiError> {
let env_var = format!("Ri_OAUTH_{}_{}", provider_name.to_uppercase(), suffix);
env::var(&env_var).map_err(|_| {
RiError::Config(format!(
"OAuth {} is not set for provider '{}'. Please set the environment variable {}",
suffix.to_lowercase(),
provider_name,
env_var
))
})
}
fn get_oauth_url(provider_name: &str, endpoint: &str) -> String {
match load_oauth_env_var(provider_name, endpoint) {
Ok(url) if !url.is_empty() => url,
_ => format!("https://{}.com/oauth/{}", provider_name, endpoint)
}
}
#[cfg(feature = "pyo3")]
use pyo3::PyResult;
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
#[derive(Debug, Clone)]
#[derive(Deserialize)]
pub struct RiAuthConfig {
pub enabled: bool,
pub jwt_secret: String,
pub jwt_expiry_secs: u64,
pub session_timeout_secs: u64,
pub oauth_providers: Vec<String>,
pub enable_api_keys: bool,
pub enable_session_auth: bool,
#[cfg(feature = "cache")]
pub oauth_cache_backend_type: crate::cache::RiCacheBackendType,
#[cfg(feature = "cache")]
pub oauth_cache_redis_url: String,
pub rate_limit_max_login_attempts: u32,
pub rate_limit_lockout_secs: u64,
pub rate_limit_window_secs: u64,
}
impl Default for RiAuthConfig {
fn default() -> Self {
Self {
enabled: true,
jwt_secret: load_jwt_secret_from_env(),
jwt_expiry_secs: 3600,
session_timeout_secs: 86400,
oauth_providers: vec![],
enable_api_keys: true,
enable_session_auth: true,
#[cfg(feature = "cache")]
oauth_cache_backend_type: crate::cache::RiCacheBackendType::Memory,
#[cfg(feature = "cache")]
oauth_cache_redis_url: "redis://127.0.0.1:6379".to_string(),
rate_limit_max_login_attempts: 5,
rate_limit_lockout_secs: 300,
rate_limit_window_secs: 900,
}
}
}
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct RiAuthModule {
config: RiAuthConfig,
jwt_manager: Arc<RiJWTManager>,
session_manager: Arc<RwLock<RiSessionManager>>,
permission_manager: Arc<RwLock<RiPermissionManager>>,
oauth_manager: Arc<RwLock<RiOAuthManager>>,
revocation_list: Arc<RiJWTRevocationList>,
rate_limiter: RiRateLimiter,
audit_logger: RiAuditLogger,
}
impl RiAuthModule {
pub async fn new(config: RiAuthConfig) -> crate::core::error::RiResult<Self> {
let jwt_manager = Arc::new(RiJWTManager::create(config.jwt_secret.clone(), config.jwt_expiry_secs));
let session_manager = Arc::new(RwLock::new(RiSessionManager::new(config.session_timeout_secs)));
let permission_manager = Arc::new(RwLock::new(RiPermissionManager::new()));
#[cfg(feature = "cache")]
let cache: Arc<dyn crate::cache::RiCache> = match config.oauth_cache_backend_type {
crate::cache::RiCacheBackendType::Memory => {
Arc::new(crate::cache::RiMemoryCache::new())
}
crate::cache::RiCacheBackendType::Redis => {
let cache = crate::cache::RiRedisCache::new(&config.oauth_cache_redis_url).await
.map_err(|e| crate::core::error::RiError::RedisError(format!("Failed to create Redis cache for OAuth: {}", e)))?;
Arc::new(cache)
}
_ => Arc::new(crate::cache::RiMemoryCache::new()),
};
#[cfg(not(feature = "cache"))]
let cache = Arc::new(crate::cache::RiMemoryCache::new());
let oauth_manager = Arc::new(RwLock::new(RiOAuthManager::new(cache)));
let revocation_list = Arc::new(RiJWTRevocationList::new());
let rate_limiter = RiRateLimiter::new(
config.rate_limit_max_login_attempts,
config.rate_limit_lockout_secs,
config.rate_limit_window_secs,
);
let audit_logger = RiAuditLogger::new(10000);
Ok(Self {
config,
jwt_manager,
session_manager,
permission_manager,
oauth_manager,
revocation_list,
rate_limiter,
audit_logger,
})
}
pub fn with_config(config: RiAuthConfig) -> Self {
let jwt_manager = Arc::new(RiJWTManager::create(config.jwt_secret.clone(), config.jwt_expiry_secs));
let session_manager = Arc::new(RwLock::new(RiSessionManager::new(config.session_timeout_secs)));
let permission_manager = Arc::new(RwLock::new(RiPermissionManager::new()));
let cache = Arc::new(crate::cache::RiMemoryCache::new());
let oauth_manager = Arc::new(RwLock::new(RiOAuthManager::new(cache)));
let revocation_list = Arc::new(RiJWTRevocationList::new());
let rate_limiter = RiRateLimiter::new(
config.rate_limit_max_login_attempts,
config.rate_limit_lockout_secs,
config.rate_limit_window_secs,
);
let audit_logger = RiAuditLogger::new(10000);
Self {
config,
jwt_manager,
session_manager,
permission_manager,
oauth_manager,
revocation_list,
rate_limiter,
audit_logger,
}
}
pub async fn new_async(config: RiAuthConfig) -> crate::core::error::RiResult<Self> {
let jwt_manager = Arc::new(RiJWTManager::create(config.jwt_secret.clone(), config.jwt_expiry_secs));
let session_manager = Arc::new(RwLock::new(RiSessionManager::new(config.session_timeout_secs)));
let permission_manager = Arc::new(RwLock::new(RiPermissionManager::new_async().await));
#[cfg(feature = "cache")]
let cache: Arc<dyn crate::cache::RiCache> = match config.oauth_cache_backend_type {
crate::cache::RiCacheBackendType::Memory => {
Arc::new(crate::cache::RiMemoryCache::new())
}
crate::cache::RiCacheBackendType::Redis => {
let cache = crate::cache::RiRedisCache::new(&config.oauth_cache_redis_url).await
.map_err(|e| crate::core::error::RiError::RedisError(format!("Failed to create Redis cache: {}", e)))?;
Arc::new(cache)
}
_ => Arc::new(crate::cache::RiMemoryCache::new()),
};
#[cfg(not(feature = "cache"))]
let cache = Arc::new(crate::cache::RiMemoryCache::new());
let oauth_manager = Arc::new(RwLock::new(RiOAuthManager::new(cache)));
let revocation_list = Arc::new(RiJWTRevocationList::new());
let rate_limiter = RiRateLimiter::new(
config.rate_limit_max_login_attempts,
config.rate_limit_lockout_secs,
config.rate_limit_window_secs,
);
let audit_logger = RiAuditLogger::new(10000);
Ok(Self {
config,
jwt_manager,
session_manager,
permission_manager,
oauth_manager,
revocation_list,
rate_limiter,
audit_logger,
})
}
pub fn revocation_list(&self) -> Arc<RiJWTRevocationList> {
self.revocation_list.clone()
}
pub async fn check_rate_limit(&self, identifier: &str) -> Result<(), u64> {
self.rate_limiter.check_and_record(identifier).await
}
pub async fn reset_rate_limit(&self, identifier: &str) {
self.rate_limiter.reset(identifier).await;
}
pub async fn log_audit_event(&self, event: RiAuditEvent) {
self.audit_logger.log(event).await;
}
pub async fn log_login_attempt(&self, username: &str, ip_address: Option<&str>, success: bool, reason: Option<&str>) {
let mut event = RiAuditEvent::new("LOGIN_ATTEMPT", if success { "login_success" } else { "login_failure" });
event = event.with_user(username);
if let Some(ip) = ip_address {
event = event.with_ip(ip);
}
if let Some(r) = reason {
event = event.with_details(r);
}
event = event.with_success(success);
self.audit_logger.log(event).await;
}
pub async fn get_audit_events(&self) -> Vec<RiAuditEvent> {
self.audit_logger.get_events().await
}
pub fn jwt_manager(&self) -> Arc<RiJWTManager> {
self.jwt_manager.clone()
}
pub fn session_manager(&self) -> Arc<RwLock<RiSessionManager>> {
self.session_manager.clone()
}
pub fn permission_manager(&self) -> Arc<RwLock<RiPermissionManager>> {
self.permission_manager.clone()
}
pub fn oauth_manager(&self) -> Arc<RwLock<RiOAuthManager>> {
self.oauth_manager.clone()
}
}
#[cfg(feature = "pyo3")]
#[pyo3::prelude::pymethods]
impl RiAuthConfig {
#[new]
#[pyo3(signature = (
enabled = true,
jwt_secret = "",
jwt_expiry_secs = 3600,
session_timeout_secs = 86400,
oauth_providers = vec![],
enable_api_keys = true,
enable_session_auth = true,
oauth_cache_backend_type = None,
oauth_cache_redis_url = "redis://127.0.0.1:6379",
rate_limit_max_login_attempts = 5,
rate_limit_lockout_secs = 300,
rate_limit_window_secs = 900
))]
fn py_new(
enabled: bool,
jwt_secret: &str,
jwt_expiry_secs: u64,
session_timeout_secs: u64,
oauth_providers: Vec<String>,
enable_api_keys: bool,
enable_session_auth: bool,
oauth_cache_backend_type: Option<String>,
oauth_cache_redis_url: &str,
rate_limit_max_login_attempts: u32,
rate_limit_lockout_secs: u64,
rate_limit_window_secs: u64,
) -> Self {
let secret = if jwt_secret.is_empty() {
load_jwt_secret_from_env()
} else {
jwt_secret.to_string()
};
#[cfg(feature = "cache")]
{
let backend_type = match oauth_cache_backend_type.as_deref() {
Some("Redis") => crate::cache::RiCacheBackendType::Redis,
_ => crate::cache::RiCacheBackendType::Memory,
};
Self {
enabled,
jwt_secret: secret,
jwt_expiry_secs,
session_timeout_secs,
oauth_providers,
enable_api_keys,
enable_session_auth,
oauth_cache_backend_type: backend_type,
oauth_cache_redis_url: oauth_cache_redis_url.to_string(),
rate_limit_max_login_attempts,
rate_limit_lockout_secs,
rate_limit_window_secs,
}
}
#[cfg(not(feature = "cache"))]
{
let _ = oauth_cache_backend_type;
let _ = oauth_cache_redis_url;
Self {
enabled,
jwt_secret: secret,
jwt_expiry_secs,
session_timeout_secs,
oauth_providers,
enable_api_keys,
enable_session_auth,
rate_limit_max_login_attempts,
rate_limit_lockout_secs,
rate_limit_window_secs,
}
}
}
#[staticmethod]
fn default() -> Self {
<Self as Default>::default()
}
#[staticmethod]
fn from_env() -> Self {
Self {
jwt_secret: load_jwt_secret_from_env(),
..Self::default()
}
}
#[getter]
fn get_enabled(&self) -> bool {
self.enabled
}
#[getter]
fn get_jwt_secret(&self) -> String {
self.jwt_secret.clone()
}
#[getter]
fn get_jwt_expiry_secs(&self) -> u64 {
self.jwt_expiry_secs
}
#[getter]
fn get_session_timeout_secs(&self) -> u64 {
self.session_timeout_secs
}
#[getter]
fn get_oauth_providers(&self) -> Vec<String> {
self.oauth_providers.clone()
}
#[getter]
fn get_enable_api_keys(&self) -> bool {
self.enable_api_keys
}
#[getter]
fn get_enable_session_auth(&self) -> bool {
self.enable_session_auth
}
}
#[cfg(feature = "pyo3")]
#[pyo3::prelude::pymethods]
impl RiAuthModule {
#[new]
fn py_new(config: RiAuthConfig) -> PyResult<Self> {
let rt = Handle::current();
rt.block_on(async {
Self::new(config).await
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
})
}
#[getter]
fn get_config(&self) -> RiAuthConfig {
self.config.clone()
}
#[getter]
fn get_jwt_expiry_secs(&self) -> u64 {
self.jwt_manager.get_token_expiry()
}
#[getter]
fn get_session_timeout_secs(&self) -> u64 {
self.config.session_timeout_secs
}
#[getter]
fn is_enabled(&self) -> bool {
self.config.enabled
}
#[getter]
fn is_api_keys_enabled(&self) -> bool {
self.config.enable_api_keys
}
#[getter]
fn is_session_auth_enabled(&self) -> bool {
self.config.enable_session_auth
}
#[getter]
fn get_oauth_providers(&self) -> Vec<String> {
self.config.oauth_providers.clone()
}
fn validate_jwt_token(&self, token: &str) -> bool {
self.jwt_manager.validate_token(token).is_ok()
}
fn generate_test_token(&self, subject: &str, roles: Vec<String>, permissions: Vec<String>) -> PyResult<String> {
self.jwt_manager.generate_token(subject, roles, permissions)
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
}
}
impl crate::core::ServiceModule for RiAuthModule {
fn name(&self) -> &str {
"Ri.Auth"
}
fn is_critical(&self) -> bool {
false
}
fn priority(&self) -> i32 {
20
}
fn dependencies(&self) -> Vec<&str> {
vec![]
}
fn init(&mut self, _ctx: &mut crate::core::RiServiceContext) -> crate::core::RiResult<()> {
Ok(())
}
fn start(&mut self, _ctx: &mut crate::core::RiServiceContext) -> crate::core::RiResult<()> {
Ok(())
}
fn shutdown(&mut self, _ctx: &mut crate::core::RiServiceContext) -> crate::core::RiResult<()> {
Ok(())
}
}
#[async_trait::async_trait]
impl crate::core::RiModule for RiAuthModule {
fn name(&self) -> &str {
"Ri.Auth"
}
fn is_critical(&self) -> bool {
false }
async fn init(&mut self, ctx: &mut RiServiceContext) -> RiResult<()> {
log::info!("Initializing Ri Auth Module");
let binding = ctx.config();
let cfg = binding.config();
if let Some(auth_config) = cfg.get("auth") {
self.config = serde_yaml::from_str(auth_config)
.unwrap_or_else(|_| RiAuthConfig::default());
}
self.jwt_manager = Arc::new(RiJWTManager::create(self.config.jwt_secret.clone(), self.config.jwt_expiry_secs));
if !self.config.oauth_providers.is_empty() {
for provider_name in &self.config.oauth_providers {
let client_id = load_oauth_env_var(provider_name, "CLIENT_ID")?;
let client_secret = load_oauth_env_var(provider_name, "CLIENT_SECRET")?;
let provider_config = crate::auth::oauth::RiOAuthProvider {
id: provider_name.clone(),
name: provider_name.clone(),
client_id,
client_secret,
auth_url: get_oauth_url(provider_name, "authorize"),
token_url: get_oauth_url(provider_name, "token"),
user_info_url: get_oauth_url(provider_name, "userinfo"),
scopes: vec!["openid".to_string(), "profile".to_string(), "email".to_string()],
enabled: true,
redirect_uri: None,
allowed_redirect_uris: vec![],
};
let oauth_mgr = self.oauth_manager.write().await;
oauth_mgr.register_provider(provider_config).await?;
log::info!("OAuth provider registered: {provider_name}");
}
}
log::info!("Ri Auth Module initialized successfully");
Ok(())
}
async fn after_shutdown(&mut self, _ctx: &mut RiServiceContext) -> RiResult<()> {
log::info!("Cleaning up Ri Auth Module");
let session_mgr = self.session_manager.write().await;
session_mgr.cleanup_all().await?;
log::info!("Ri Auth Module cleanup completed");
Ok(())
}
}