ant_quic/relay/

session_manager.rs

//! Session management for relay connections with complete state machine.

use crate::relay::{
    AuthToken, RelayAuthenticator, RelayConnection, RelayConnectionConfig, RelayError, RelayResult,
};
use ed25519_dalek::VerifyingKey;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::sync::mpsc;

/// Unique session identifier
pub type SessionId = u32;

/// Session configuration
#[derive(Debug, Clone)]
pub struct SessionConfig {
    /// Maximum number of concurrent sessions
    pub max_sessions: usize,
    /// Default session timeout
    pub default_timeout: Duration,
    /// Session cleanup interval
    pub cleanup_interval: Duration,
    /// Default bandwidth limit per session
    pub default_bandwidth_limit: u64,
}

impl Default for SessionConfig {
    fn default() -> Self {
        Self {
            max_sessions: 100,
            default_timeout: Duration::from_secs(300), // 5 minutes
            cleanup_interval: Duration::from_secs(30), // 30 seconds
            default_bandwidth_limit: 1048576,          // 1 MB/s
        }
    }
}

/// Session state in the relay state machine
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SessionState {
    /// Session requested but not yet established
    Pending,
    /// Session active and forwarding data
    Active,
    /// Session terminating gracefully
    Terminating,
    /// Session terminated
    Terminated,
    /// Session failed due to error
    Failed { reason: String },
}

/// Information about a relay session
#[derive(Debug, Clone)]
pub struct RelaySessionInfo {
    /// Session identifier
    pub session_id: SessionId,
    /// Client address
    pub client_addr: SocketAddr,
    /// Target peer connection ID
    pub peer_connection_id: Vec<u8>,
    /// Current session state
    pub state: SessionState,
    /// Session creation time
    pub created_at: Instant,
    /// Last activity time
    pub last_activity: Instant,
    /// Bandwidth limit
    pub bandwidth_limit: u64,
    /// Session timeout
    pub timeout: Duration,
    /// Bytes transferred
    pub bytes_sent: u64,
    pub bytes_received: u64,
}

/// Session manager for handling relay connections
#[derive(Debug)]
pub struct SessionManager {
    /// Configuration
    config: SessionConfig,
    /// Active sessions
    sessions: Arc<Mutex<HashMap<SessionId, RelaySessionInfo>>>,
    /// Active connections
    connections: Arc<Mutex<HashMap<SessionId, Arc<RelayConnection>>>>,
    /// Authenticator for token verification
    authenticator: RelayAuthenticator,
    /// Trusted peer keys for authentication
    trusted_keys: Arc<Mutex<HashMap<SocketAddr, VerifyingKey>>>,
    /// Next session ID
    next_session_id: Arc<Mutex<SessionId>>,
    /// Event channels
    event_sender: mpsc::UnboundedSender<SessionEvent>,
    /// Last cleanup time
    last_cleanup: Arc<Mutex<Instant>>,
}

/// Events generated by session management
#[derive(Debug, Clone)]
pub enum SessionEvent {
    /// New session requested
    SessionRequested {
        session_id: SessionId,
        client_addr: SocketAddr,
        peer_connection_id: Vec<u8>,
        auth_token: AuthToken,
    },
    /// Session established successfully
    SessionEstablished {
        session_id: SessionId,
        client_addr: SocketAddr,
    },
    /// Session terminated
    SessionTerminated {
        session_id: SessionId,
        reason: String,
    },
    /// Session failed
    SessionFailed {
        session_id: SessionId,
        error: RelayError,
    },
    /// Data forwarded through session
    DataForwarded {
        session_id: SessionId,
        bytes: usize,
        direction: ForwardDirection,
    },
}

/// Direction of data forwarding
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ForwardDirection {
    /// From client to peer
    ClientToPeer,
    /// From peer to client
    PeerToClient,
}

impl SessionManager {
    /// Create a new session manager
    pub fn new(config: SessionConfig) -> (Self, mpsc::UnboundedReceiver<SessionEvent>) {
        let (event_sender, event_receiver) = mpsc::unbounded_channel();

        let manager = Self {
            config,
            sessions: Arc::new(Mutex::new(HashMap::new())),
            connections: Arc::new(Mutex::new(HashMap::new())),
            authenticator: RelayAuthenticator::new(),
            trusted_keys: Arc::new(Mutex::new(HashMap::new())),
            next_session_id: Arc::new(Mutex::new(1)),
            event_sender,
            last_cleanup: Arc::new(Mutex::new(Instant::now())),
        };

        (manager, event_receiver)
    }

    /// Add a trusted peer key for authentication
    pub fn add_trusted_key(&self, addr: SocketAddr, key: VerifyingKey) {
        let mut trusted_keys = self.trusted_keys.lock().unwrap();
        trusted_keys.insert(addr, key);
    }

    /// Remove a trusted peer key
    pub fn remove_trusted_key(&self, addr: &SocketAddr) {
        let mut trusted_keys = self.trusted_keys.lock().unwrap();
        trusted_keys.remove(addr);
    }

    /// Generate next session ID
    fn next_session_id(&self) -> SessionId {
        let mut next_id = self.next_session_id.lock().unwrap();
        let id = *next_id;
        *next_id = next_id.wrapping_add(1);
        if *next_id == 0 {
            *next_id = 1; // Skip 0 as invalid session ID
        }
        id
    }

    /// Request a new relay session
    pub fn request_session(
        &self,
        client_addr: SocketAddr,
        peer_connection_id: Vec<u8>,
        auth_token: AuthToken,
    ) -> RelayResult<SessionId> {
        // Check session limit
        {
            let sessions = self.sessions.lock().unwrap();
            if sessions.len() >= self.config.max_sessions {
                return Err(RelayError::ResourceExhausted {
                    resource_type: "sessions".to_string(),
                    current_usage: sessions.len() as u64,
                    limit: self.config.max_sessions as u64,
                });
            }
        }

        // Verify authentication token
        let trusted_keys = self.trusted_keys.lock().unwrap();
        let peer_key =
            trusted_keys
                .get(&client_addr)
                .ok_or_else(|| RelayError::AuthenticationFailed {
                    reason: format!("No trusted key for address {}", client_addr),
                })?;

        self.authenticator.verify_token(&auth_token, peer_key)?;

        // Generate session ID
        let session_id = self.next_session_id();

        // Create session info
        let now = Instant::now();
        let session_info = RelaySessionInfo {
            session_id,
            client_addr,
            peer_connection_id: peer_connection_id.clone(),
            state: SessionState::Pending,
            created_at: now,
            last_activity: now,
            bandwidth_limit: auth_token.bandwidth_limit as u64,
            timeout: Duration::from_secs(auth_token.timeout_seconds as u64),
            bytes_sent: 0,
            bytes_received: 0,
        };

        // Store session
        {
            let mut sessions = self.sessions.lock().unwrap();
            sessions.insert(session_id, session_info);
        }

        // Send event
        let _ = self.event_sender.send(SessionEvent::SessionRequested {
            session_id,
            client_addr,
            peer_connection_id,
            auth_token,
        });

        Ok(session_id)
    }

    /// Establish a relay session
    pub fn establish_session(&self, session_id: SessionId) -> RelayResult<()> {
        let (client_addr, bandwidth_limit) = {
            let mut sessions = self.sessions.lock().unwrap();
            let session = sessions
                .get_mut(&session_id)
                .ok_or(RelayError::SessionError {
                    session_id: Some(session_id),
                    kind: crate::relay::error::SessionErrorKind::NotFound,
                })?;

            if session.state != SessionState::Pending {
                return Err(RelayError::SessionError {
                    session_id: Some(session_id),
                    kind: crate::relay::error::SessionErrorKind::InvalidState {
                        current_state: format!("{:?}", session.state),
                        expected_state: "Pending".to_string(),
                    },
                });
            }

            session.state = SessionState::Active;
            session.last_activity = Instant::now();

            (session.client_addr, session.bandwidth_limit)
        };

        // Create relay connection
        let (event_tx, _event_rx) = mpsc::unbounded_channel();
        let (_action_tx, action_rx) = mpsc::unbounded_channel();

        let mut connection_config = RelayConnectionConfig::default();
        connection_config.bandwidth_limit = bandwidth_limit;

        let connection = RelayConnection::new(
            session_id,
            client_addr,
            connection_config,
            event_tx,
            action_rx,
        );

        // Store connection
        {
            let mut connections = self.connections.lock().unwrap();
            connections.insert(session_id, Arc::new(connection));
        }

        // Send event
        let _ = self.event_sender.send(SessionEvent::SessionEstablished {
            session_id,
            client_addr,
        });

        Ok(())
    }

    /// Terminate a relay session
    pub fn terminate_session(&self, session_id: SessionId, reason: String) -> RelayResult<()> {
        // Update session state
        {
            let mut sessions = self.sessions.lock().unwrap();
            if let Some(session) = sessions.get_mut(&session_id) {
                session.state = SessionState::Terminated;
                session.last_activity = Instant::now();
            }
        }

        // Terminate connection
        {
            let mut connections = self.connections.lock().unwrap();
            if let Some(connection) = connections.remove(&session_id) {
                let _ = connection.terminate(reason.clone());
            }
        }

        // Send event
        let _ = self
            .event_sender
            .send(SessionEvent::SessionTerminated { session_id, reason });

        Ok(())
    }

    /// Forward data through a relay session
    pub fn forward_data(
        &self,
        session_id: SessionId,
        data: Vec<u8>,
        direction: ForwardDirection,
    ) -> RelayResult<()> {
        let connection = {
            let connections = self.connections.lock().unwrap();
            connections
                .get(&session_id)
                .cloned()
                .ok_or(RelayError::SessionError {
                    session_id: Some(session_id),
                    kind: crate::relay::error::SessionErrorKind::NotFound,
                })?
        };

        // Forward data based on direction
        match direction {
            ForwardDirection::ClientToPeer => {
                connection.send_data(data.clone())?;
            }
            ForwardDirection::PeerToClient => {
                connection.receive_data(data.clone())?;
            }
        }

        // Update session statistics
        {
            let mut sessions = self.sessions.lock().unwrap();
            if let Some(session) = sessions.get_mut(&session_id) {
                session.last_activity = Instant::now();
                match direction {
                    ForwardDirection::ClientToPeer => {
                        session.bytes_sent += data.len() as u64;
                    }
                    ForwardDirection::PeerToClient => {
                        session.bytes_received += data.len() as u64;
                    }
                }
            }
        }

        // Send event
        let _ = self.event_sender.send(SessionEvent::DataForwarded {
            session_id,
            bytes: data.len(),
            direction,
        });

        Ok(())
    }

    /// Get session information
    pub fn get_session(&self, session_id: SessionId) -> Option<RelaySessionInfo> {
        let sessions = self.sessions.lock().unwrap();
        sessions.get(&session_id).cloned()
    }

    /// List all active sessions
    pub fn list_sessions(&self) -> Vec<RelaySessionInfo> {
        let sessions = self.sessions.lock().unwrap();
        sessions.values().cloned().collect()
    }

    /// Get session count
    pub fn session_count(&self) -> usize {
        let sessions = self.sessions.lock().unwrap();
        sessions.len()
    }

    /// Clean up expired sessions
    pub fn cleanup_expired_sessions(&self) -> RelayResult<usize> {
        let mut last_cleanup = self.last_cleanup.lock().unwrap();
        let now = Instant::now();

        // Only cleanup if enough time has passed
        if now.duration_since(*last_cleanup) < self.config.cleanup_interval {
            return Ok(0);
        }

        *last_cleanup = now;
        drop(last_cleanup);

        let mut expired_sessions = Vec::new();

        // Find expired sessions
        {
            let sessions = self.sessions.lock().unwrap();
            for (session_id, session_info) in sessions.iter() {
                let age = now.duration_since(session_info.last_activity);
                if age > session_info.timeout {
                    expired_sessions.push(*session_id);
                }
            }
        }

        // Remove expired sessions
        let cleanup_count = expired_sessions.len();
        for session_id in expired_sessions {
            let _ = self.terminate_session(session_id, "Session expired".to_string());

            // Remove from sessions map
            let mut sessions = self.sessions.lock().unwrap();
            sessions.remove(&session_id);
        }

        Ok(cleanup_count)
    }

    /// Get session manager statistics
    pub fn get_statistics(&self) -> SessionManagerStats {
        let sessions = self.sessions.lock().unwrap();
        let connections = self.connections.lock().unwrap();

        let mut active_count = 0;
        let mut pending_count = 0;
        let mut total_bytes_sent = 0;
        let mut total_bytes_received = 0;

        for session in sessions.values() {
            match session.state {
                SessionState::Active => active_count += 1,
                SessionState::Pending => pending_count += 1,
                _ => {}
            }
            total_bytes_sent += session.bytes_sent;
            total_bytes_received += session.bytes_received;
        }

        SessionManagerStats {
            total_sessions: sessions.len(),
            active_sessions: active_count,
            pending_sessions: pending_count,
            total_connections: connections.len(),
            total_bytes_sent,
            total_bytes_received,
        }
    }
}

/// Session manager statistics
#[derive(Debug, Clone)]
pub struct SessionManagerStats {
    pub total_sessions: usize,
    pub active_sessions: usize,
    pub pending_sessions: usize,
    pub total_connections: usize,
    pub total_bytes_sent: u64,
    pub total_bytes_received: u64,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::relay::AuthToken;
    use ed25519_dalek::SigningKey;
    use rand::rngs::OsRng;
    use std::net::{IpAddr, Ipv4Addr};

    fn test_addr() -> SocketAddr {
        SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080)
    }

    #[test]
    fn test_session_manager_creation() {
        let config = SessionConfig::default();
        let (manager, _event_rx) = SessionManager::new(config);

        let stats = manager.get_statistics();
        assert_eq!(stats.total_sessions, 0);
        assert_eq!(stats.active_sessions, 0);
    }

    #[test]
    fn test_trusted_key_management() {
        let config = SessionConfig::default();
        let (manager, _event_rx) = SessionManager::new(config);

        let signing_key = SigningKey::generate(&mut OsRng);
        let verifying_key = signing_key.verifying_key();
        let addr = test_addr();

        manager.add_trusted_key(addr, verifying_key);

        // Should be able to create a session with trusted key
        let auth_token = AuthToken::new(1024, 300, &signing_key).unwrap();
        let result = manager.request_session(addr, vec![1, 2, 3], auth_token);
        assert!(result.is_ok());

        // Remove trusted key
        manager.remove_trusted_key(&addr);

        // Should fail without trusted key
        let auth_token2 = AuthToken::new(1024, 300, &signing_key).unwrap();
        let result2 = manager.request_session(addr, vec![4, 5, 6], auth_token2);
        assert!(result2.is_err());
    }

    #[test]
    fn test_session_request_and_establishment() {
        let config = SessionConfig::default();
        let (manager, _event_rx) = SessionManager::new(config);

        let signing_key = SigningKey::generate(&mut OsRng);
        let verifying_key = signing_key.verifying_key();
        let addr = test_addr();

        manager.add_trusted_key(addr, verifying_key);

        // Request session
        let auth_token = AuthToken::new(1024, 300, &signing_key).unwrap();
        let session_id = manager
            .request_session(addr, vec![1, 2, 3], auth_token)
            .unwrap();

        // Check session exists and is pending
        let session = manager.get_session(session_id).unwrap();
        assert_eq!(session.state, SessionState::Pending);
        assert_eq!(session.client_addr, addr);

        // Establish session
        assert!(manager.establish_session(session_id).is_ok());

        // Check session is now active
        let session = manager.get_session(session_id).unwrap();
        assert_eq!(session.state, SessionState::Active);
    }

    #[test]
    fn test_session_limit() {
        let mut config = SessionConfig::default();
        config.max_sessions = 2;
        let (manager, _event_rx) = SessionManager::new(config);

        let signing_key = SigningKey::generate(&mut OsRng);
        let verifying_key = signing_key.verifying_key();
        let addr = test_addr();

        manager.add_trusted_key(addr, verifying_key);

        // Create maximum sessions
        for i in 0..2 {
            let auth_token = AuthToken::new(1024, 300, &signing_key).unwrap();
            let result = manager.request_session(addr, vec![i], auth_token);
            assert!(result.is_ok());
        }

        // Third session should fail
        let auth_token = AuthToken::new(1024, 300, &signing_key).unwrap();
        let result = manager.request_session(addr, vec![3], auth_token);
        assert!(result.is_err());
    }

    #[test]
    fn test_session_termination() {
        let config = SessionConfig::default();
        let (manager, _event_rx) = SessionManager::new(config);

        let signing_key = SigningKey::generate(&mut OsRng);
        let verifying_key = signing_key.verifying_key();
        let addr = test_addr();

        manager.add_trusted_key(addr, verifying_key);

        // Create and establish session
        let auth_token = AuthToken::new(1024, 300, &signing_key).unwrap();
        let session_id = manager
            .request_session(addr, vec![1, 2, 3], auth_token)
            .unwrap();
        manager.establish_session(session_id).unwrap();

        // Terminate session
        let reason = "Test termination".to_string();
        assert!(manager.terminate_session(session_id, reason).is_ok());

        // Check session is terminated
        let session = manager.get_session(session_id).unwrap();
        assert_eq!(session.state, SessionState::Terminated);
    }

    #[test]
    fn test_data_forwarding() {
        let config = SessionConfig::default();
        let (manager, _event_rx) = SessionManager::new(config);

        let signing_key = SigningKey::generate(&mut OsRng);
        let verifying_key = signing_key.verifying_key();
        let addr = test_addr();

        manager.add_trusted_key(addr, verifying_key);

        // Create and establish session
        let auth_token = AuthToken::new(1024, 300, &signing_key).unwrap();
        let session_id = manager
            .request_session(addr, vec![1, 2, 3], auth_token)
            .unwrap();
        manager.establish_session(session_id).unwrap();

        // Forward data
        let data = vec![1, 2, 3, 4, 5];
        assert!(
            manager
                .forward_data(session_id, data.clone(), ForwardDirection::ClientToPeer)
                .is_ok()
        );
        assert!(
            manager
                .forward_data(session_id, data, ForwardDirection::PeerToClient)
                .is_ok()
        );

        // Check statistics updated
        let session = manager.get_session(session_id).unwrap();
        assert_eq!(session.bytes_sent, 5);
        assert_eq!(session.bytes_received, 5);
    }

    #[test]
    fn test_session_cleanup() {
        let mut config = SessionConfig::default();
        config.cleanup_interval = Duration::from_millis(1);
        let (manager, _event_rx) = SessionManager::new(config);

        let signing_key = SigningKey::generate(&mut OsRng);
        let verifying_key = signing_key.verifying_key();
        let addr = test_addr();

        manager.add_trusted_key(addr, verifying_key);

        // Create session with very short timeout
        let auth_token = AuthToken::new(1024, 1, &signing_key).unwrap(); // 1 second timeout
        let _session_id = manager
            .request_session(addr, vec![1, 2, 3], auth_token)
            .unwrap();

        assert_eq!(manager.session_count(), 1);

        // Wait for session to expire
        std::thread::sleep(Duration::from_millis(2));

        // Cleanup should remove expired session
        let cleanup_count = manager.cleanup_expired_sessions().unwrap();
        assert!(cleanup_count > 0);
    }

    #[test]
    fn test_session_id_generation() {
        let config = SessionConfig::default();
        let (manager, _event_rx) = SessionManager::new(config);

        let signing_key = SigningKey::generate(&mut OsRng);
        let verifying_key = signing_key.verifying_key();
        let addr = test_addr();

        manager.add_trusted_key(addr, verifying_key);

        // Generate multiple session IDs
        let mut session_ids = Vec::new();
        for i in 0..10 {
            let auth_token = AuthToken::new(1024, 300, &signing_key).unwrap();
            let session_id = manager.request_session(addr, vec![i], auth_token).unwrap();
            session_ids.push(session_id);
        }

        // All IDs should be unique and non-zero
        for id in &session_ids {
            assert!(*id != 0);
        }

        let unique_ids: std::collections::HashSet<_> = session_ids.iter().collect();
        assert_eq!(unique_ids.len(), session_ids.len());
    }
}