frag_datagram/

server.rs

/*
 * Copyright (c) Peter Bjorklund. All rights reserved. https://github.com/piot/frag-datagram
 * Licensed under the MIT License. See LICENSE in the project root for license information.
 */
use crate::{
    FLAG_CONNECT_REQUEST, FLAG_CONNECT_RESPONSE, Receiver, Sender, read_connect_request,
    read_datagram, write_connect_response,
};
use std::collections::{HashMap, HashSet};
use std::time::Duration;
use tracing::trace;

/// A complete bidirectional connection that combines Sender + Receiver
pub struct ServerConnection {
    connection_id: u16,
    address_hash: u64,
    sender: Sender,
    receiver: Receiver,
}

impl ServerConnection {
    /// Create a new server connection
    #[must_use]
    pub fn new(connection_id: u16, address_hash: u64, max_fragments: usize) -> Self {
        Self {
            connection_id,
            address_hash,
            sender: Sender::new(connection_id),
            receiver: Receiver::new(max_fragments, address_hash),
        }
    }

    /// Send data to this connection's client
    pub fn send(&mut self, payload: &[u8]) -> Vec<Vec<u8>> {
        self.sender.send(payload)
    }

    /// Process incoming datagram for this connection
    pub fn receive(
        &mut self,
        datagram: &[u8],
        sender_address_hash: u64,
    ) -> Option<(Vec<u8>, bool)> {
        self.receiver.receive(datagram, sender_address_hash)
    }

    /// Check if this connection is inactive
    #[must_use]
    pub fn is_inactive(&self, timeout: Duration) -> bool {
        self.receiver.is_inactive(timeout)
    }

    /// Get connection ID
    #[must_use]
    pub const fn connection_id(&self) -> u16 {
        self.connection_id
    }

    /// Get address hash
    #[must_use]
    pub const fn address_hash(&self) -> u64 {
        self.address_hash
    }
}

/// Server-side hub that manages complete bidirectional connections
pub struct ServerHub {
    connections: HashMap<u16, ServerConnection>,
    seen_connection_requests: HashSet<(u64, u32)>,
    max_fragments_per_connection: usize,
    max_connections: usize,
    connection_timeout: Duration,
}

impl ServerHub {
    /// Create a new server hub
    #[must_use]
    pub fn new(
        max_fragments_per_connection: usize,
        max_connections: usize,
        connection_timeout_seconds: u64,
    ) -> Self {
        Self {
            connections: HashMap::new(),
            seen_connection_requests: HashSet::new(),
            max_fragments_per_connection,
            max_connections,
            connection_timeout: Duration::from_secs(connection_timeout_seconds),
        }
    }

    /// Process incoming datagram
    /// Returns:
    /// - `Some((connection_id`, payload, `response_datagrams`, `is_new_connection`)) for complete messages
    /// - None for incomplete/invalid datagrams
    pub fn receive(
        &mut self,
        datagram: &[u8],
        address_hash: u64,
    ) -> Option<(u16, Vec<u8>, Vec<Vec<u8>>, bool)> {
        // Parse header to get connection_id and flags
        let (header, _) = read_datagram(datagram)?;

        // Handle connection requests directly
        if header.flags == FLAG_CONNECT_REQUEST {
            trace!(len = datagram.len(), "got connect request");

            let (_, connect_req, user_payload) = read_connect_request(datagram)?;
            let request_key = (address_hash, connect_req.request_id);
            let is_new_request = !self.seen_connection_requests.contains(&request_key);

            if is_new_request {
                self.seen_connection_requests.insert(request_key);
            }

            let (connection_id, is_new) = if is_new_request {
                (self.create_connection(address_hash)?, true)
            } else {
                let existing_conn_id = self
                    .connections
                    .iter()
                    .find(|(_, connection)| connection.address_hash() == address_hash)
                    .map(|(&id, _)| id)?;
                (existing_conn_id, false)
            };

            if let Some(response) =
                self.create_connection_response(connect_req.request_id, connection_id, &[])
            {
                // Return the user payload from this fragment so the caller can reconstruct
                return Some((connection_id, user_payload.to_vec(), vec![response], is_new));
            }

            return None;
        }

        if header.flags == FLAG_CONNECT_RESPONSE {
            return None;
        }

        let connection_id = header.connection_id;
        if let Some(connection) = self.connections.get_mut(&connection_id)
            && let Some((payload, _is_new_conn)) = connection.receive(datagram, address_hash)
        {
            return Some((connection_id, payload, vec![], false));
        }

        None
    }

    /// Create a new connection after accepting a connection request
    pub fn create_connection(&mut self, address_hash: u64) -> Option<u16> {
        if self.connections.len() >= self.max_connections {
            return None;
        }

        // Generate a unique connection ID
        let mut connection_id = 1u16;
        while self.connections.contains_key(&connection_id) {
            connection_id = connection_id.wrapping_add(1);
            if connection_id == 0 {
                connection_id = 1; // Skip 0, reserved for connection requests
            }
        }

        let connection = ServerConnection::new(
            connection_id,
            address_hash,
            self.max_fragments_per_connection,
        );
        self.connections.insert(connection_id, connection);

        Some(connection_id)
    }

    /// Send data to a specific connection
    pub fn send_to(&mut self, connection_id: u16, payload: &[u8]) -> Option<Vec<Vec<u8>>> {
        self.connections
            .get_mut(&connection_id)
            .map(|conn| conn.send(payload))
    }

    /// Create a connection response datagram
    #[must_use]
    pub fn create_connection_response(
        &self,
        request_id: u32,
        connection_id: u16,
        payload: &[u8],
    ) -> Option<Vec<u8>> {
        let mut buf = vec![0u8; 1024];
        if let Some(len) =
            write_connect_response(&mut buf, request_id, connection_id, 0, 0, payload)
        {
            buf.truncate(len);
            Some(buf)
        } else {
            None
        }
    }

    /// Remove inactive connections
    pub fn cleanup_inactive_connections(&mut self) -> Vec<u16> {
        let timeout = self.connection_timeout;
        let mut inactive_connections = Vec::new();

        for (&connection_id, connection) in &self.connections {
            if connection.is_inactive(timeout) {
                inactive_connections.push(connection_id);
            }
        }

        for connection_id in &inactive_connections {
            self.connections.remove(connection_id);
        }

        inactive_connections
    }

    /// Get connection count
    #[must_use]
    pub fn connection_count(&self) -> usize {
        self.connections.len()
    }

    /// Check if connection exists
    #[must_use]
    pub fn has_connection(&self, connection_id: u16) -> bool {
        self.connections.contains_key(&connection_id)
    }

    /// Remove a specific connection
    pub fn remove_connection(&mut self, connection_id: u16) -> bool {
        self.connections.remove(&connection_id).is_some()
    }
}

/// Utility functions for generating address hashes from network addresses
pub mod address_hash {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};

    /// Generate a hash from IPv4 address and port
    #[must_use]
    pub fn hash_ipv4(ip: u32, port: u16) -> u64 {
        let mut hasher = DefaultHasher::new();
        ip.hash(&mut hasher);
        port.hash(&mut hasher);
        hasher.finish()
    }

    /// Generate a hash from IPv6 address and port
    #[must_use]
    pub fn hash_ipv6(ip: [u8; 16], port: u16) -> u64 {
        let mut hasher = DefaultHasher::new();
        ip.hash(&mut hasher);
        port.hash(&mut hasher);
        hasher.finish()
    }

    /// Generate a hash from a string representation of an address
    #[must_use]
    pub fn hash_from_string(address_string: &str) -> u64 {
        let mut hasher = DefaultHasher::new();
        address_string.hash(&mut hasher);
        hasher.finish()
    }
}