corevpn_auth/

session.rs

//! Authentication Session Management

use std::collections::HashMap;
use std::time::Duration;

use chrono::{DateTime, Utc};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::{AuthError, AuthState, Result, TokenSet, UserInfo};

/// Authentication session
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuthSession {
    /// Session ID
    pub id: String,
    /// Authentication state (for OAuth2 flow)
    pub auth_state: Option<AuthState>,
    /// Token set (after successful authentication)
    pub tokens: Option<TokenSet>,
    /// User information
    pub user_info: Option<UserInfo>,
    /// Provider name
    pub provider: String,
    /// Session creation time
    pub created_at: DateTime<Utc>,
    /// Session expiration time
    pub expires_at: DateTime<Utc>,
    /// Last activity time
    pub last_activity: DateTime<Utc>,
    /// Associated VPN session ID (if connected)
    pub vpn_session_id: Option<String>,
    /// IP address of the client
    pub client_ip: Option<String>,
    /// Additional metadata
    pub metadata: HashMap<String, String>,
}

impl AuthSession {
    /// Create a new auth session
    pub fn new(provider: &str, lifetime: Duration) -> Self {
        let now = Utc::now();
        let auth_state = AuthState::new(Duration::from_secs(600)); // 10 min for OAuth flow

        Self {
            id: Uuid::new_v4().to_string(),
            auth_state: Some(auth_state),
            tokens: None,
            user_info: None,
            provider: provider.to_string(),
            created_at: now,
            expires_at: now + chrono::Duration::from_std(lifetime).unwrap(),
            last_activity: now,
            vpn_session_id: None,
            client_ip: None,
            metadata: HashMap::new(),
        }
    }

    /// Check if session is expired
    pub fn is_expired(&self) -> bool {
        Utc::now() > self.expires_at
    }

    /// Check if session is authenticated
    pub fn is_authenticated(&self) -> bool {
        self.tokens.is_some() && self.user_info.is_some()
    }

    /// Check if tokens need refresh
    pub fn needs_token_refresh(&self) -> bool {
        if let Some(tokens) = &self.tokens {
            tokens.expires_within(chrono::Duration::minutes(5))
        } else {
            false
        }
    }

    /// Update tokens
    pub fn update_tokens(&mut self, tokens: TokenSet) {
        self.tokens = Some(tokens);
        self.last_activity = Utc::now();
    }

    /// Update user info
    pub fn update_user_info(&mut self, user_info: UserInfo) {
        self.user_info = Some(user_info);
        self.last_activity = Utc::now();
    }

    /// Mark authentication complete
    pub fn complete_auth(&mut self, tokens: TokenSet, user_info: UserInfo) {
        self.tokens = Some(tokens);
        self.user_info = Some(user_info);
        self.auth_state = None; // Clear auth state after successful auth
        self.last_activity = Utc::now();
    }

    /// Associate with VPN session
    pub fn associate_vpn_session(&mut self, vpn_session_id: &str) {
        self.vpn_session_id = Some(vpn_session_id.to_string());
        self.last_activity = Utc::now();
    }

    /// Extend session lifetime
    pub fn extend(&mut self, duration: Duration) {
        self.expires_at = Utc::now() + chrono::Duration::from_std(duration).unwrap();
    }

    /// Touch session (update last activity)
    pub fn touch(&mut self) {
        self.last_activity = Utc::now();
    }

    /// Get session duration
    pub fn duration(&self) -> chrono::Duration {
        Utc::now() - self.created_at
    }

    /// Get idle time
    pub fn idle_time(&self) -> chrono::Duration {
        Utc::now() - self.last_activity
    }

    /// Get the OAuth2 state value
    pub fn state(&self) -> Option<&str> {
        self.auth_state.as_ref().map(|s| s.state.as_str())
    }

    /// Get user email
    pub fn email(&self) -> Option<&str> {
        self.user_info.as_ref().and_then(|u| u.email.as_deref())
    }

    /// Get user display name
    pub fn display_name(&self) -> Option<&str> {
        self.user_info.as_ref().and_then(|u| u.name.as_deref())
    }
}

/// Authentication session manager
pub struct AuthSessionManager {
    /// Sessions by ID
    sessions: RwLock<HashMap<String, AuthSession>>,
    /// Sessions by OAuth2 state
    sessions_by_state: RwLock<HashMap<String, String>>,
    /// Default session lifetime
    default_lifetime: Duration,
    /// Maximum sessions per user
    max_sessions_per_user: usize,
}

impl AuthSessionManager {
    /// Create a new session manager
    pub fn new(default_lifetime: Duration, max_sessions_per_user: usize) -> Self {
        Self {
            sessions: RwLock::new(HashMap::new()),
            sessions_by_state: RwLock::new(HashMap::new()),
            default_lifetime,
            max_sessions_per_user,
        }
    }

    /// Create a new session
    pub fn create_session(&self, provider: &str) -> AuthSession {
        let session = AuthSession::new(provider, self.default_lifetime);

        // Store session
        let mut sessions = self.sessions.write();
        let mut by_state = self.sessions_by_state.write();

        if let Some(state) = session.state() {
            by_state.insert(state.to_string(), session.id.clone());
        }
        sessions.insert(session.id.clone(), session.clone());

        session
    }

    /// Get session by ID
    pub fn get_session(&self, id: &str) -> Option<AuthSession> {
        self.sessions.read().get(id).cloned()
    }

    /// Get session by OAuth2 state
    pub fn get_session_by_state(&self, state: &str) -> Option<AuthSession> {
        let session_id = self.sessions_by_state.read().get(state)?.clone();
        self.get_session(&session_id)
    }

    /// Update session
    pub fn update_session(&self, session: &AuthSession) -> Result<()> {
        let mut sessions = self.sessions.write();
        if sessions.contains_key(&session.id) {
            sessions.insert(session.id.clone(), session.clone());
            Ok(())
        } else {
            Err(AuthError::SessionNotFound)
        }
    }

    /// Remove session
    pub fn remove_session(&self, id: &str) -> Option<AuthSession> {
        let mut sessions = self.sessions.write();
        let mut by_state = self.sessions_by_state.write();

        if let Some(session) = sessions.remove(id) {
            if let Some(state) = session.state() {
                by_state.remove(state);
            }
            Some(session)
        } else {
            None
        }
    }

    /// Get all sessions for a user (by email)
    pub fn get_user_sessions(&self, email: &str) -> Vec<AuthSession> {
        self.sessions
            .read()
            .values()
            .filter(|s| s.email() == Some(email))
            .cloned()
            .collect()
    }

    /// Remove all sessions for a user
    pub fn remove_user_sessions(&self, email: &str) -> usize {
        let mut sessions = self.sessions.write();
        let mut by_state = self.sessions_by_state.write();

        let to_remove: Vec<_> = sessions
            .iter()
            .filter(|(_, s)| s.email() == Some(email))
            .map(|(id, s)| (id.clone(), s.state().map(String::from)))
            .collect();

        for (id, state) in &to_remove {
            sessions.remove(id);
            if let Some(s) = state {
                by_state.remove(s);
            }
        }

        to_remove.len()
    }

    /// Cleanup expired sessions
    pub fn cleanup_expired(&self) -> usize {
        let mut sessions = self.sessions.write();
        let mut by_state = self.sessions_by_state.write();

        let before = sessions.len();

        let expired: Vec<_> = sessions
            .iter()
            .filter(|(_, s)| s.is_expired())
            .map(|(id, s)| (id.clone(), s.state().map(String::from)))
            .collect();

        for (id, state) in &expired {
            sessions.remove(id);
            if let Some(s) = state {
                by_state.remove(s);
            }
        }

        before - sessions.len()
    }

    /// Get session count
    pub fn session_count(&self) -> usize {
        self.sessions.read().len()
    }

    /// Get all active sessions
    pub fn active_sessions(&self) -> Vec<AuthSession> {
        self.sessions
            .read()
            .values()
            .filter(|s| !s.is_expired() && s.is_authenticated())
            .cloned()
            .collect()
    }
}

impl Default for AuthSessionManager {
    fn default() -> Self {
        Self::new(Duration::from_secs(86400), 5) // 24 hours, 5 sessions per user
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_auth_session() {
        let mut session = AuthSession::new("google", Duration::from_secs(3600));

        assert!(!session.is_expired());
        assert!(!session.is_authenticated());
        assert!(session.state().is_some());
    }

    #[test]
    fn test_session_manager() {
        let manager = AuthSessionManager::default();

        let session = manager.create_session("google");
        let state = session.state().unwrap().to_string();

        // Get by ID
        let found = manager.get_session(&session.id);
        assert!(found.is_some());

        // Get by state
        let found = manager.get_session_by_state(&state);
        assert!(found.is_some());

        // Remove
        manager.remove_session(&session.id);
        assert!(manager.get_session(&session.id).is_none());
    }

    #[test]
    fn test_session_lifecycle() {
        let manager = AuthSessionManager::default();

        let mut session = manager.create_session("google");

        // Simulate successful auth
        let tokens = TokenSet {
            access_token: "test-token".to_string(),
            refresh_token: Some("refresh-token".to_string()),
            id_token: None,
            expires_at: Utc::now() + chrono::Duration::hours(1),
            token_type: "Bearer".to_string(),
            scopes: vec![],
        };

        let user_info = UserInfo {
            sub: "user123".to_string(),
            email: Some("user@example.com".to_string()),
            email_verified: true,
            name: Some("Test User".to_string()),
            given_name: None,
            family_name: None,
            picture: None,
            groups: vec![],
            provider: "google".to_string(),
        };

        session.complete_auth(tokens, user_info);

        assert!(session.is_authenticated());
        assert!(session.auth_state.is_none()); // Cleared after auth
        assert_eq!(session.email(), Some("user@example.com"));
    }
}