ax-net 0.13.1

//! Public socket facade.
//!
//! This module defines the protocol-independent socket API used by syscall
//! layers: common send/recv flags, extended address families, shutdown modes,
//! and the `SocketOps` trait implemented by TCP, UDP, raw, Unix, and vsock
//! transports.
//!
//! # Compatibility Boundary
//!
//! The syscall layer should not need to know whether a socket is backed by
//! smoltcp, an in-kernel Unix transport, or a vsock connection manager. It
//! passes `SocketAddrEx`, `SendOptions`, and `RecvOptions` into this facade, and
//! each concrete transport maps them onto its own semantics.
//!
//! # Design Rule
//!
//! This module contains dispatch and common ABI shapes only. Protocol behavior
//! such as TCP accept queues, UDP corking, raw packet format, or Unix ancillary
//! data delivery belongs in the corresponding transport module.

use alloc::{boxed::Box, vec::Vec};
use core::{
    any::Any,
    fmt::{self, Debug},
    net::SocketAddr,
    time::Duration,
};

use ax_io::prelude::*;
use axpoll::{ExclusiveRegistrationSink, IoEvents, Pollable, SharedRegistrationSink};
use bitflags::bitflags;
use enum_dispatch::enum_dispatch;

#[cfg(feature = "vsock")]
use crate::vsock::{VsockAddr, VsockSocket};
use crate::{
    NetError, NetResult,
    options::{Configurable, GetSocketOption, SetSocketOption, UnixCredentials},
    raw::RawSocket,
    tcp::TcpSocket,
    udp::UdpSocket,
    unix::{UnixSocket, UnixSocketAddr},
};

/// Extended socket address supporting IP, Unix, and vsock address families.
#[derive(Clone, Debug)]
pub enum SocketAddrEx {
    /// An IP (v4/v6) socket address.
    Ip(SocketAddr),
    /// A Unix domain socket address.
    Unix(UnixSocketAddr),
    /// A vsock socket address.
    #[cfg(feature = "vsock")]
    Vsock(VsockAddr),
}

impl SocketAddrEx {
    /// Convert into an IP socket address, or return an error if not IP.
    pub fn into_ip(self) -> NetResult<SocketAddr> {
        match self {
            SocketAddrEx::Ip(addr) => Ok(addr),
            SocketAddrEx::Unix(_) => Err(NetError::AddressFamilyUnsupported),
            #[cfg(feature = "vsock")]
            SocketAddrEx::Vsock(_) => Err(NetError::AddressFamilyUnsupported),
        }
    }

    /// Convert into a Unix socket address, or return an error if not Unix.
    pub fn into_unix(self) -> NetResult<UnixSocketAddr> {
        match self {
            SocketAddrEx::Unix(addr) => Ok(addr),
            SocketAddrEx::Ip(_) => Err(NetError::AddressFamilyUnsupported),
            #[cfg(feature = "vsock")]
            SocketAddrEx::Vsock(_) => Err(NetError::AddressFamilyUnsupported),
        }
    }

    /// Convert into a vsock address, or return an error if not vsock.
    #[cfg(feature = "vsock")]
    pub fn into_vsock(self) -> NetResult<VsockAddr> {
        match self {
            SocketAddrEx::Ip(_) => Err(NetError::AddressFamilyUnsupported),
            SocketAddrEx::Unix(_) => Err(NetError::AddressFamilyUnsupported),
            SocketAddrEx::Vsock(addr) => Ok(addr),
        }
    }
}

bitflags! {
    /// Flags for sending data to a socket.
    ///
    /// These values match Linux MSG_* constants so that `from_bits_retain(flags)`
    /// from the syscall layer preserves the correct flags.
    ///
    /// See [`SocketOps::send`].
    #[derive(Default, Debug, Clone, Copy)]
    pub struct SendFlags: u32 {
        /// Sends out-of-band data on sockets that support it (e.g. SOCK_STREAM).
        const OOB = 0x01;
        /// Don't use a gateway to send the packet, send to hosts only on
        /// directly connected networks.
        const DONTROUTE = 0x04;
        /// Enables nonblocking operation; if the operation would block,
        /// `EAGAIN` or `EWOULDBLOCK` is returned.
        const DONTWAIT = 0x40;
        /// Terminates a record (SOCK_SEQPACKET).
        const EOR = 0x80;
        /// Sends only if a connection confirm is pending (UDP/RAW, Linux specific).
        const CONFIRM = 0x800;
        /// Requests not to send SIGPIPE on errors on stream oriented sockets
        /// when the other end breaks the connection.
        const NOSIGNAL = 0x4000;
        /// More data will be sent; used to cork/coalesce sends (UDP/TCP).
        const MORE = 0x8000;
    }
}

bitflags! {
    /// Flags for receiving data from a socket.
    ///
    /// See [`SocketOps::recv`].
    #[derive(Default, Debug, Clone, Copy)]
    pub struct RecvFlags: u32 {
        /// Receive data without removing it from the queue.
        const PEEK = 0x01;
        /// For datagram-like sockets, requires [`SocketOps::recv`] to return
        /// the real size of the datagram, even when it is larger than the
        /// buffer.
        const TRUNCATE = 0x02;
        /// Requests out-of-band data (`MSG_OOB`). Only stream sockets that
        /// support urgent data honor it; others reject it with `EOPNOTSUPP`.
        const OOB = 0x04;
        /// Per-call non-blocking override (`MSG_DONTWAIT`). Does NOT
        /// change the socket's own `O_NONBLOCK` state.
        const DONTWAIT = 0x40;
    }
}

/// Ancillary control message payload, carried opaquely so the socket layer
/// stays protocol-independent. Cloneable so `recvmsg(MSG_PEEK)` can duplicate
/// the ancillary data without consuming the record: SCM_RIGHTS fds are cloned
/// (sharing the open file description), matching Linux `unix_peek_fds` /
/// `scm_fp_dup`.
pub trait CMsgPayload: Any + Send + Sync {
    /// Duplicate into a fresh payload for peek delivery.
    fn clone_box(&self) -> Box<dyn CMsgPayload>;
    /// Recover a `Box<dyn Any>` for owned downcast at the syscall layer.
    fn into_any(self: Box<Self>) -> Box<dyn Any + Send + Sync>;
}
impl<T: Any + Send + Sync + Clone> CMsgPayload for T {
    fn clone_box(&self) -> Box<dyn CMsgPayload> {
        Box::new(self.clone())
    }
    fn into_any(self: Box<Self>) -> Box<dyn Any + Send + Sync> {
        self
    }
}
// Opaque payload; mirror the std `dyn Any` Debug impl so containers deriving
// Debug still compile.
impl core::fmt::Debug for dyn CMsgPayload {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.write_str("CMsgPayload { .. }")
    }
}

/// Type alias for ancillary control message data.
pub type CMsgData = Box<dyn CMsgPayload>;

/// IP ancillary data reported through `recvmsg`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IpCmsg {
    /// IPv4 hop limit for `IP_RECVTTL`.
    Ipv4Ttl(u8),
    /// IPv4 TOS byte for `IP_RECVTOS`.
    Ipv4Tos(u8),
    /// IPv6 traffic-class byte for `IPV6_RECVTCLASS`.
    Ipv6TrafficClass(u8),
}

/// Transport-independent socket-level ancillary data reported through
/// `recvmsg`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SocketCmsg {
    /// Sender credentials requested with `SO_PASSCRED`.
    Credentials(UnixCredentials),
    /// Wall-clock receive timestamp requested with `SO_TIMESTAMP`.
    Timestamp(Duration),
}

/// Options for sending data to a socket.
///
/// See [`SocketOps::send`].
#[derive(Default, Debug)]
pub struct SendOptions {
    /// Destination address for the message.
    pub to: Option<SocketAddrEx>,
    /// Send flags.
    pub flags: SendFlags,
    /// Ancillary control messages.
    pub cmsg: Vec<CMsgData>,
    /// Real credentials of the task performing this send operation.
    pub sender_credentials: Option<UnixCredentials>,
}

/// Options for receiving data from a socket.
///
/// See [`SocketOps::recv`].
#[derive(Default)]
pub struct RecvOptions<'a> {
    /// If set, the sender's address is written here.
    pub from: Option<&'a mut SocketAddrEx>,
    /// Receive flags.
    pub flags: RecvFlags,
    /// If set, ancillary control messages are appended here.
    pub cmsg: Option<&'a mut Vec<CMsgData>>,
    /// If set and the datagram was truncated, this is set to `true`.
    pub truncated: Option<&'a mut bool>,
}
impl Debug for RecvOptions<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("RecvOptions")
            .field("from", &self.from)
            .field("flags", &self.flags)
            .finish()
    }
}

/// Kind of shutdown operation to perform on a socket.
#[derive(Debug, Clone, Copy)]
pub enum Shutdown {
    /// Shut down the read half.
    Read,
    /// Shut down the write half.
    Write,
    /// Shut down both halves.
    Both,
}

/// Progress of a connection attempt whose completion is readiness-driven.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ConnectStatus {
    /// The transport completed the connection synchronously or asynchronously.
    Connected,
    /// The transport started an asynchronous connection and is not ready yet.
    InProgress,
}

/// OS-facing policy for driving one task-neutral socket future.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct SocketWaitPolicy {
    /// Whether this operation must return after the registration recheck.
    pub nonblocking: bool,
    /// Optional relative timeout owned and enforced by the consuming OS.
    pub timeout: Option<Duration>,
}

fn socket_wait_policy(
    socket: &(impl Configurable + ?Sized),
    send: bool,
    extra_nonblocking: bool,
) -> NetResult<SocketWaitPolicy> {
    let mut nonblocking = false;
    socket.get_option(GetSocketOption::NonBlocking(&mut nonblocking))?;
    let mut timeout = Duration::ZERO;
    if send {
        socket.get_option(GetSocketOption::SendTimeout(&mut timeout))?;
    } else {
        socket.get_option(GetSocketOption::ReceiveTimeout(&mut timeout))?;
    }
    Ok(SocketWaitPolicy {
        nonblocking: nonblocking || extra_nonblocking,
        timeout: (!timeout.is_zero()).then_some(timeout),
    })
}
impl Shutdown {
    /// Returns `true` if the read half should be shut down.
    pub fn has_read(&self) -> bool {
        matches!(self, Shutdown::Read | Shutdown::Both)
    }

    /// Returns `true` if the write half should be shut down.
    pub fn has_write(&self) -> bool {
        matches!(self, Shutdown::Write | Shutdown::Both)
    }
}

/// Operations that can be performed on a socket.
#[enum_dispatch]
pub trait SocketOps: Configurable {
    /// Binds an unbound socket to the given address and port.
    fn bind(&self, local_addr: SocketAddrEx) -> NetResult;
    /// Starts a connection without waiting for asynchronous completion.
    fn start_connect(&self, remote_addr: SocketAddrEx) -> NetResult<ConnectStatus>;
    /// Checks a previously started asynchronous connection.
    fn connect_status(&self) -> NetResult<ConnectStatus> {
        Ok(ConnectStatus::Connected)
    }

    /// Starts listening on the bound address and port.
    fn listen(&self, _backlog: usize) -> NetResult {
        Err(NetError::OperationNotSupported)
    }
    /// Returns whether this socket currently accepts incoming connections.
    fn is_listening(&self) -> bool {
        false
    }
    /// Attempts to accept one connection without parking the caller.
    fn try_accept(&self) -> NetResult<Socket> {
        Err(NetError::OperationNotSupported)
    }

    /// Attempts to send data without parking the caller.
    fn try_send(&self, src: impl Read + IoBuf, options: &mut SendOptions) -> NetResult<usize>;
    /// Attempts to receive data without parking the caller.
    fn try_recv(
        &self,
        dst: impl Write + IoBufMut,
        options: &mut RecvOptions<'_>,
    ) -> NetResult<usize>;
    /// Returns the number of bytes that can be read without blocking.
    fn recv_available(&self) -> NetResult<usize> {
        Err(NetError::OperationNotSupported)
    }

    /// Get the local endpoint of the socket.
    fn local_addr(&self) -> NetResult<SocketAddrEx>;
    /// Get the remote endpoint of the socket.
    fn peer_addr(&self) -> NetResult<SocketAddrEx>;

    /// Shutdown the socket, closing the connection.
    fn shutdown(&self, how: Shutdown) -> NetResult;

    /// Returns the send wait policy without acquiring scheduler ownership.
    fn send_wait_policy(&self, extra_nonblocking: bool) -> NetResult<SocketWaitPolicy> {
        socket_wait_policy(self, true, extra_nonblocking)
    }

    /// Returns the receive wait policy without acquiring scheduler ownership.
    fn recv_wait_policy(&self, extra_nonblocking: bool) -> NetResult<SocketWaitPolicy> {
        socket_wait_policy(self, false, extra_nonblocking)
    }
}

impl<T: SocketOps + ?Sized> SocketOps for Box<T> {
    fn bind(&self, local_addr: SocketAddrEx) -> NetResult {
        (**self).bind(local_addr)
    }

    fn start_connect(&self, remote_addr: SocketAddrEx) -> NetResult<ConnectStatus> {
        (**self).start_connect(remote_addr)
    }

    fn connect_status(&self) -> NetResult<ConnectStatus> {
        (**self).connect_status()
    }

    fn listen(&self, backlog: usize) -> NetResult {
        (**self).listen(backlog)
    }

    fn is_listening(&self) -> bool {
        (**self).is_listening()
    }

    fn try_accept(&self) -> NetResult<Socket> {
        (**self).try_accept()
    }

    fn try_send(&self, src: impl Read + IoBuf, options: &mut SendOptions) -> NetResult<usize> {
        (**self).try_send(src, options)
    }

    fn try_recv(
        &self,
        dst: impl Write + IoBufMut,
        options: &mut RecvOptions<'_>,
    ) -> NetResult<usize> {
        (**self).try_recv(dst, options)
    }

    fn recv_available(&self) -> NetResult<usize> {
        (**self).recv_available()
    }

    fn local_addr(&self) -> NetResult<SocketAddrEx> {
        (**self).local_addr()
    }

    fn peer_addr(&self) -> NetResult<SocketAddrEx> {
        (**self).peer_addr()
    }

    fn shutdown(&self, how: Shutdown) -> NetResult {
        (**self).shutdown(how)
    }
}

/// Network socket abstraction.
#[enum_dispatch(Configurable, SocketOps)]
pub enum Socket {
    /// UDP socket.
    Udp(Box<UdpSocket>),
    /// TCP socket.
    Tcp(Box<TcpSocket>),
    /// Raw IP socket.
    Raw(Box<RawSocket>),
    /// Unix domain socket.
    Unix(Box<UnixSocket>),
    /// Virtio socket.
    #[cfg(feature = "vsock")]
    Vsock(Box<VsockSocket>),
}

impl From<UdpSocket> for Socket {
    fn from(socket: UdpSocket) -> Self {
        Self::Udp(Box::new(socket))
    }
}

impl From<TcpSocket> for Socket {
    fn from(socket: TcpSocket) -> Self {
        Self::Tcp(Box::new(socket))
    }
}

impl From<UnixSocket> for Socket {
    fn from(socket: UnixSocket) -> Self {
        Self::Unix(Box::new(socket))
    }
}

#[cfg(feature = "vsock")]
impl From<VsockSocket> for Socket {
    fn from(socket: VsockSocket) -> Self {
        Self::Vsock(Box::new(socket))
    }
}

impl Pollable for Socket {
    fn poll(&self) -> IoEvents {
        match self {
            Socket::Tcp(tcp) => tcp.poll(),
            Socket::Udp(udp) => udp.poll(),
            Socket::Raw(raw) => raw.poll(),
            Socket::Unix(unix) => unix.poll(),
            #[cfg(feature = "vsock")]
            Socket::Vsock(vsock) => vsock.poll(),
        }
    }

    unsafe fn register_shared(&self, sink: &mut dyn SharedRegistrationSink, events: IoEvents) {
        match self {
            Socket::Tcp(tcp) => unsafe { tcp.register_shared(sink, events) },
            Socket::Udp(udp) => unsafe { udp.register_shared(sink, events) },
            Socket::Raw(raw) => unsafe { raw.register_shared(sink, events) },
            Socket::Unix(unix) => unsafe { unix.register_shared(sink, events) },
            #[cfg(feature = "vsock")]
            Socket::Vsock(vsock) => unsafe { vsock.register_shared(sink, events) },
        }
    }

    unsafe fn register_exclusive(
        &self,
        sink: &mut dyn ExclusiveRegistrationSink,
        events: IoEvents,
    ) {
        match self {
            Socket::Tcp(tcp) => unsafe { tcp.register_exclusive(sink, events) },
            Socket::Udp(udp) => unsafe { udp.register_exclusive(sink, events) },
            Socket::Raw(raw) => unsafe { raw.register_exclusive(sink, events) },
            Socket::Unix(unix) => unsafe { unix.register_exclusive(sink, events) },
            #[cfg(feature = "vsock")]
            Socket::Vsock(vsock) => unsafe { vsock.register_exclusive(sink, events) },
        }
    }
}