aerosocket_client/

connection.rs

//! WebSocket client connection handling for AeroSocket
//!
//! This module provides connection management for WebSocket clients.

use aerosocket_core::prelude::Bytes;
use aerosocket_core::{Message, Result};
use std::net::SocketAddr;

/// Represents a WebSocket client connection
#[derive(Debug)]
pub struct ClientConnection {
    /// Server address
    remote_addr: SocketAddr,
    /// Connection state
    state: ConnectionState,
    /// Connection metadata
    metadata: ConnectionMetadata,
}

/// Connection state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConnectionState {
    /// Connection is being established
    Connecting,
    /// Connection is established and ready
    Connected,
    /// Connection is closing
    Closing,
    /// Connection is closed
    Closed,
}

/// Connection metadata
#[derive(Debug, Clone)]
pub struct ConnectionMetadata {
    /// WebSocket subprotocol
    pub subprotocol: Option<String>,
    /// WebSocket extensions
    pub extensions: Vec<String>,
    /// Connection established time
    pub established_at: std::time::Instant,
    /// Last activity time
    pub last_activity_at: std::time::Instant,
    /// Messages sent count
    pub messages_sent: u64,
    /// Messages received count
    pub messages_received: u64,
    /// Bytes sent count
    pub bytes_sent: u64,
    /// Bytes received count
    pub bytes_received: u64,
}

impl ClientConnection {
    /// Create a new client connection
    pub fn new(remote_addr: SocketAddr) -> Self {
        Self {
            remote_addr,
            state: ConnectionState::Connecting,
            metadata: ConnectionMetadata {
                subprotocol: None,
                extensions: Vec::new(),
                established_at: std::time::Instant::now(),
                last_activity_at: std::time::Instant::now(),
                messages_sent: 0,
                messages_received: 0,
                bytes_sent: 0,
                bytes_received: 0,
            },
        }
    }

    /// Get the remote address
    pub fn remote_addr(&self) -> SocketAddr {
        self.remote_addr
    }

    /// Get the connection state
    pub fn state(&self) -> ConnectionState {
        self.state
    }

    /// Get the connection metadata
    pub fn metadata(&self) -> &ConnectionMetadata {
        &self.metadata
    }

    /// Send a message
    pub async fn send(&mut self, message: Message) -> Result<()> {
        // TODO: Implement actual message sending
        self.metadata.messages_sent += 1;
        let message_len = match &message {
            Message::Text(text) => text.len(),
            Message::Binary(data) => data.len(),
            Message::Ping(data) => data.len(),
            Message::Pong(data) => data.len(),
            Message::Close(close_msg) => close_msg.len(),
        };
        self.metadata.bytes_sent += message_len as u64;
        self.metadata.last_activity_at = std::time::Instant::now();
        Ok(())
    }

    /// Send a text message
    pub async fn send_text(&mut self, text: impl AsRef<str>) -> Result<()> {
        let message = Message::text(text.as_ref().to_string());
        self.send(message).await
    }

    /// Send a binary message
    pub async fn send_binary(&mut self, data: impl Into<Bytes>) -> Result<()> {
        let message = Message::binary(data);
        self.send(message).await
    }

    /// Send a ping message
    pub async fn ping(&mut self, data: Option<&[u8]>) -> Result<()> {
        let message = Message::ping(data.map(|d| d.to_vec()));
        self.send(message).await
    }

    /// Send a pong message
    pub async fn pong(&mut self, data: Option<&[u8]>) -> Result<()> {
        let message = Message::pong(data.map(|d| d.to_vec()));
        self.send(message).await
    }

    /// Receive the next message
    pub async fn next(&mut self) -> Result<Option<Message>> {
        // TODO: Implement actual message receiving
        Ok(None)
    }

    /// Close the connection
    pub async fn close(&mut self, code: Option<u16>, reason: Option<&str>) -> Result<()> {
        self.state = ConnectionState::Closing;
        let message = Message::close(code, reason.map(|s| s.to_string()));
        self.send(message).await?;
        self.state = ConnectionState::Closed;
        Ok(())
    }

    /// Check if the connection is established
    pub fn is_connected(&self) -> bool {
        self.state == ConnectionState::Connected
    }

    /// Check if the connection is closed
    pub fn is_closed(&self) -> bool {
        self.state == ConnectionState::Closed
    }

    /// Get the connection age
    pub fn age(&self) -> std::time::Duration {
        self.metadata.established_at.elapsed()
    }

    /// Get the time since last activity
    pub fn idle_time(&self) -> std::time::Duration {
        self.metadata.last_activity_at.elapsed()
    }

    /// Set the connection as connected
    pub fn set_connected(&mut self) {
        self.state = ConnectionState::Connected;
        self.metadata.established_at = std::time::Instant::now();
        self.metadata.last_activity_at = std::time::Instant::now();
    }

    /// Set the subprotocol
    pub fn set_subprotocol(&mut self, subprotocol: String) {
        self.metadata.subprotocol = Some(subprotocol);
    }

    /// Add an extension
    pub fn add_extension(&mut self, extension: String) {
        self.metadata.extensions.push(extension);
    }
}

/// Connection handle for managing connections
#[derive(Debug)]
pub struct ClientConnectionHandle {
    /// Connection ID
    id: u64,
    /// Connection reference
    connection: std::sync::Arc<std::sync::Mutex<ClientConnection>>,
}

impl ClientConnectionHandle {
    /// Create a new connection handle
    pub fn new(id: u64, connection: ClientConnection) -> Self {
        Self {
            id,
            connection: std::sync::Arc::new(std::sync::Mutex::new(connection)),
        }
    }

    /// Get the connection ID
    pub fn id(&self) -> u64 {
        self.id
    }

    /// Get a reference to the connection
    pub fn connection(&self) -> &std::sync::Arc<std::sync::Mutex<ClientConnection>> {
        &self.connection
    }

    /// Try to lock the connection
    pub fn try_lock(&self) -> aerosocket_core::Result<std::sync::MutexGuard<'_, ClientConnection>> {
        self.connection
            .try_lock()
            .map_err(|e| aerosocket_core::Error::Other(format!("Poison error: {}", e)))
    }
}

impl Clone for ClientConnectionHandle {
    fn clone(&self) -> Self {
        Self {
            id: self.id,
            connection: self.connection.clone(),
        }
    }
}

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

    #[test]
    fn test_client_connection_creation() {
        let remote = "127.0.0.1:8080".parse().unwrap();
        let conn = ClientConnection::new(remote);

        assert_eq!(conn.remote_addr(), remote);
        assert_eq!(conn.state(), ConnectionState::Connecting);
        assert!(!conn.is_connected());
        assert!(!conn.is_closed());
    }

    #[test]
    fn test_client_connection_handle() {
        let remote = "127.0.0.1:8080".parse().unwrap();
        let conn = ClientConnection::new(remote);
        let handle = ClientConnectionHandle::new(1, conn);

        assert_eq!(handle.id(), 1);
        assert!(handle.try_lock().is_ok());
    }

    #[test]
    fn test_connection_state_transitions() {
        let remote = "127.0.0.1:8080".parse().unwrap();
        let mut conn = ClientConnection::new(remote);

        assert_eq!(conn.state(), ConnectionState::Connecting);

        conn.set_connected();
        assert_eq!(conn.state(), ConnectionState::Connected);
        assert!(conn.is_connected());
    }

    #[tokio::test]
    async fn test_message_sending() {
        let remote = "127.0.0.1:8080".parse().unwrap();
        let mut conn = ClientConnection::new(remote);

        // These should not fail even without actual implementation
        assert!(conn.send_text("Hello").await.is_ok());
        assert!(conn.send_binary(&[1u8, 2, 3][..]).await.is_ok());
        assert!(conn.ping(None).await.is_ok());
        assert!(conn.pong(None).await.is_ok());
    }
}