ax-net 0.13.1

//! Shared smoltcp socket-set wrapper.
//!
//! ax-net keeps one global smoltcp `SocketSet` behind this wrapper. The extra
//! UDP bind table fills the per-address bind semantics that smoltcp itself does
//! not track for all POSIX cases.
//!
//! # Ownership
//!
//! All TCP, UDP, and raw smoltcp socket handles live in the same handle space.
//! This is what allows the service poller, router snooping path, listen table,
//! and orphan reaper to coordinate without per-interface socket duplication.
//!
//! # UDP Side Table
//!
//! smoltcp validates whether a UDP socket can bind, but ax-net needs
//! Linux-style wildcard/specific-address conflict checks across sockets. The
//! `udp_binds` table records only successful public binds and is cleaned when a
//! socket is removed.
//!
//! # Lock Boundary
//!
//! The wrapper lock protects smoltcp socket state. Callers should keep the lock
//! scoped to direct socket access and avoid waking tasks or acquiring the outer
//! service lock while it is held.

use alloc::{vec, vec::Vec};

use ax_sync::Mutex;
use hashbrown::HashMap;
use smoltcp::{
    iface::{SocketHandle, SocketSet},
    socket::AnySocket,
    wire::IpAddress,
};

use crate::{NetError, NetResult, addr::listen_addrs_conflict};

/// One UDP bind ownership record. Several records share a port only when every
/// binder requested SO_REUSEPORT on the identical local address, mirroring
/// Linux's reuseport group semantics.
#[derive(Clone, Debug)]
struct UdpBoundEntry {
    /// `None` represents a wildcard bind.
    addr: Option<IpAddress>,
    reuse_port: bool,
    handle: SocketHandle,
}

/// Global socket container plus protocol-specific side tables.
pub(crate) struct SocketSetWrapper<'a> {
    /// The shared smoltcp socket set.
    pub inner: Mutex<SocketSet<'a>>,
    /// UDP bind ownership tracked with Linux-style wildcard/reuseport conflicts.
    udp_binds: Mutex<HashMap<u16, Vec<UdpBoundEntry>>>,
}

impl<'a> SocketSetWrapper<'a> {
    /// Creates an empty wrapper around smoltcp's socket set.
    pub fn new() -> Self {
        Self {
            inner: Mutex::new(SocketSet::new(vec![])),
            udp_binds: Mutex::new(HashMap::new()),
        }
    }

    /// Adds a smoltcp socket and returns its global handle.
    pub fn add<T: AnySocket<'a>>(&self, socket: T) -> SocketHandle {
        let handle = self.inner.lock().add(socket);
        debug!("socket {}: created", handle);
        handle
    }

    /// Runs a closure with mutable access to one smoltcp socket.
    pub fn with_socket_mut<T: AnySocket<'a>, R, F>(&self, handle: SocketHandle, f: F) -> R
    where
        F: FnOnce(&mut T) -> R,
    {
        let mut set = self.inner.lock();
        let socket = set.get_mut(handle);
        f(socket)
    }

    /// Records a public UDP bind after checking address/reuseport conflicts.
    ///
    /// A binder joins an existing group on the same port only when it and every
    /// colliding owner requested SO_REUSEPORT on the exact same local address;
    /// any other overlap is rejected with `EADDRINUSE`.
    pub fn udp_bind(
        &self,
        handle: SocketHandle,
        addr: IpAddress,
        port: u16,
        reuse_port: bool,
    ) -> NetResult {
        if port == 0 {
            return Ok(());
        }
        let addr = (!addr.is_unspecified()).then_some(addr);
        let mut binds = self.udp_binds.lock();
        let entries = binds.entry(port).or_default();
        if entries
            .iter()
            .any(|entry| udp_binds_conflict(entry, addr, reuse_port))
        {
            return Err(NetError::AddrInUse);
        }
        entries.push(UdpBoundEntry {
            addr,
            reuse_port,
            handle,
        });
        Ok(())
    }

    /// Returns whether a UDP port can be used for an ephemeral bind.
    pub fn udp_port_available(&self, addr: IpAddress, port: u16) -> bool {
        if port == 0 {
            return true;
        }
        let addr = (!addr.is_unspecified()).then_some(addr);
        match self.udp_binds.lock().get(&port) {
            None => true,
            Some(entries) => !entries
                .iter()
                .any(|entry| listen_addrs_conflict(entry.addr, addr)),
        }
    }

    /// Removes any UDP bind table entries owned by `handle`.
    pub fn udp_unbind(&self, handle: SocketHandle) {
        self.udp_binds.lock().retain(|_, entries| {
            entries.retain(|entry| entry.handle != handle);
            !entries.is_empty()
        });
    }

    /// Removes a socket and all wrapper-maintained side-table state.
    pub fn remove(&self, handle: SocketHandle) {
        self.udp_unbind(handle);
        self.inner.lock().remove(handle);
        debug!("socket {}: destroyed", handle);
    }
}

/// A new UDP bind conflicts with an existing owner unless both requested
/// SO_REUSEPORT on the exact same local address.
fn udp_binds_conflict(entry: &UdpBoundEntry, addr: Option<IpAddress>, reuse_port: bool) -> bool {
    listen_addrs_conflict(entry.addr, addr)
        && !(reuse_port && entry.reuse_port && entry.addr == addr)
}

#[cfg(test)]
mod tests {
    use alloc::vec;

    use smoltcp::{
        iface::SocketSet,
        socket::udp,
        storage::PacketMetadata,
        wire::{IpAddress, Ipv4Address},
    };

    use super::*;

    fn addr(a: u8, b: u8, c: u8, d: u8) -> IpAddress {
        IpAddress::Ipv4(Ipv4Address::new(a, b, c, d))
    }

    fn entry(addr: Option<IpAddress>, reuse_port: bool) -> UdpBoundEntry {
        let mut sockets = SocketSet::new(vec![]);
        let handle = sockets.add(udp::Socket::new(
            udp::PacketBuffer::new(vec![PacketMetadata::EMPTY; 1], vec![0; 8]),
            udp::PacketBuffer::new(vec![PacketMetadata::EMPTY; 1], vec![0; 8]),
        ));
        UdpBoundEntry {
            handle,
            addr,
            reuse_port,
        }
    }

    #[test]
    fn wildcard_and_identical_udp_bindings_conflict() {
        let specific = Some(addr(192, 0, 2, 10));
        let owner = entry(specific, false);

        assert!(udp_binds_conflict(&owner, specific, false));
        assert!(udp_binds_conflict(&owner, None, false));
        assert!(!udp_binds_conflict(
            &owner,
            Some(addr(198, 51, 100, 20)),
            false,
        ));
    }

    #[test]
    fn udp_reuseport_requires_an_exact_reuseport_group() {
        let specific = Some(addr(127, 0, 0, 1));
        let plain = entry(specific, false);
        let reuse = entry(specific, true);

        assert!(udp_binds_conflict(&plain, specific, true));
        assert!(!udp_binds_conflict(&reuse, specific, true));
        assert!(udp_binds_conflict(&reuse, None, true));
    }
}