rust-kernel

//! BSD-style sockets: what `socket()` returns, as a file descriptor.
//!
//! A socket is a `File` (read = recv, write = send, close = close), plus
//! the socket-only calls connect/bind/listen/accept/sendto/recvfrom. Each
//! wraps one smoltcp socket (kept in the shared `SocketSet` in net/mod.rs):
//!   * TCP  (SOCK_STREAM): a reliable byte stream to one peer.
//!   * UDP  (SOCK_DGRAM): single datagrams, to/from anyone.
//!   * ICMP (SOCK_DGRAM + IPPROTO_ICMP): Linux's "ping socket": you send an
//!     ICMP echo request (8-byte header + payload), you receive the replies.
//!     The kernel fills in the identifier, so programs needn't be root.
//!
//! Blocking: smoltcp never blocks, it just says "no data yet". So e.g.
//! `recv` sleeps on `net::CHANGED` (woken after every `net::poll`) until
//! its socket has data, the peer closed, or the timeout (SO_RCVTIMEO) ran
//! out. No busy-waiting: a blocked program costs no CPU.
//!
//! A listening TCP socket is a small *backlog* of smoltcp sockets all
//! listening on the port (a smoltcp socket handles exactly one
//! connection). `accept` takes one that got connected, makes it a new fd,
//! and puts a fresh listener in its place.

use super::{Net, poll, with_net};
use crate::errno::Errno;
use crate::fs::file::{File, KIND_SOCKET, Stat};
use alloc::vec;
use alloc::vec::Vec;
use smoltcp::iface::SocketHandle;
use smoltcp::socket::{icmp, tcp, udp};
use smoltcp::wire::{IpAddress, IpEndpoint, Ipv4Address};
use spin::Mutex;
use x86_64::instructions::interrupts::without_interrupts;

const TCP_BUFFER: usize = 16 * 1024;
const DGRAM_BUFFER: usize = 8 * 1024;
const DGRAM_PACKETS: usize = 16;
const MAX_BACKLOG: usize = 4;
const CONNECT_TIMEOUT_US: u64 = 10_000_000;
/// A closed connection gives up on an unresponsive peer after this long.
const CLOSE_TIMEOUT_MS: u64 = 10_000;

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Kind {
    Tcp,
    Udp,
    Icmp,
}

struct State {
    handle: Option<SocketHandle>,
    listeners: Vec<SocketHandle>,
    local_port: u16,
    /// Default destination (UDP/ICMP after `connect`).
    remote: Option<IpEndpoint>,
    /// SO_RCVTIMEO: how long recv/accept may block (None = forever).
    timeout_us: Option<u64>,
}

pub struct Socket {
    pub kind: Kind,
    state: Mutex<State>,
}

/// An IPv4 address + port, as plain numbers (the syscall layer converts
/// from/to `struct sockaddr_in`).
pub type Addr = ([u8; 4], u16);

fn endpoint((ip, port): Addr) -> IpEndpoint {
    IpEndpoint::new(IpAddress::Ipv4(Ipv4Address::from(ip)), port)
}

fn addr(ep: IpEndpoint) -> Addr {
    #[allow(irrefutable_let_patterns)]
    let IpAddress::Ipv4(ip) = ep.addr else { return ([0; 4], ep.port) };
    (ip.octets(), ep.port)
}

fn new_tcp() -> tcp::Socket<'static> {
    tcp::Socket::new(tcp::SocketBuffer::new(vec![0; TCP_BUFFER]), tcp::SocketBuffer::new(vec![0; TCP_BUFFER]))
}

impl Socket {
    /// socket(): a new, unconnected socket.
    pub fn new(kind: Kind) -> Result<Socket, Errno> {
        if !super::present() {
            return Err(Errno::ENETDOWN);
        }
        let mut state = State { handle: None, listeners: Vec::new(), local_port: 0, remote: None, timeout_us: None };
        match kind {
            Kind::Tcp => {} // created by connect/listen
            Kind::Udp => {
                let buf = || udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY; DGRAM_PACKETS], vec![0; DGRAM_BUFFER]);
                state.handle = with_net(|n| n.sockets.add(udp::Socket::new(buf(), buf())));
            }
            Kind::Icmp => {
                let buf = || icmp::PacketBuffer::new(vec![icmp::PacketMetadata::EMPTY; DGRAM_PACKETS], vec![0; DGRAM_BUFFER]);
                state.handle = with_net(|n| {
                    let mut sock = icmp::Socket::new(buf(), buf());
                    let ident = n.ephemeral_port();
                    sock.bind(icmp::Endpoint::Ident(ident)).ok();
                    state.local_port = ident; // our echo identifier
                    n.sockets.add(sock)
                });
            }
        }
        Ok(Socket { kind, state: Mutex::new(state) })
    }

    /// Lock our state and the network (always in that order, interrupts
    /// off) and run `f`.
    fn locked<R>(&self, f: impl FnOnce(&mut State, &mut Net) -> R) -> Result<R, Errno> {
        without_interrupts(|| {
            let mut st = self.state.lock();
            with_net(|n| f(&mut st, n)).ok_or(Errno::ENETDOWN)
        })
    }

    /// Block until `f` returns Some (re-checked after every network poll),
    /// or the deadline passes (-> EAGAIN, or `on_timeout`).
    fn wait<R>(
        &self,
        timeout_us: Option<u64>,
        mut f: impl FnMut(&mut State, &mut Net) -> Option<Result<R, Errno>>,
    ) -> Result<R, Errno> {
        let deadline = timeout_us.map_or(u64::MAX, crate::time::deadline_after_us);
        super::CHANGED
            .wait_until_deadline(deadline, || match self.locked(&mut f) {
                Ok(r) => r,
                Err(e) => Some(Err(e)),
            })
            .unwrap_or(Err(Errno::EAGAIN))
    }

    fn recv_timeout(&self) -> Option<u64> {
        without_interrupts(|| self.state.lock().timeout_us)
    }

    pub fn set_recv_timeout(&self, us: Option<u64>) {
        without_interrupts(|| self.state.lock().timeout_us = us);
    }

    pub fn bind(&self, port: u16) -> Result<(), Errno> {
        self.locked(|st, n| {
            if st.local_port != 0 && self.kind != Kind::Icmp {
                return Err(Errno::EINVAL); // already bound
            }
            match self.kind {
                Kind::Tcp => st.local_port = port,
                Kind::Udp => {
                    let port = if port == 0 { n.ephemeral_port() } else { port };
                    n.sockets.get_mut::<udp::Socket>(st.handle.unwrap()).bind(port).map_err(|_| Errno::EADDRINUSE)?;
                    st.local_port = port;
                }
                Kind::Icmp => {} // ICMP has no ports; we keep our identifier
            }
            Ok(())
        })?
    }

    /// connect(): TCP does the three-way handshake (SYN, SYN-ACK, ACK) and
    /// waits for it; UDP/ICMP just remember the default destination.
    pub fn connect(&self, to: Addr) -> Result<(), Errno> {
        super::wait_configured(5_000_000)?;
        if self.kind != Kind::Tcp {
            if self.kind == Kind::Udp && without_interrupts(|| self.state.lock().local_port) == 0 {
                self.bind(0)?;
            }
            return self.locked(|st, _| st.remote = Some(endpoint(to)));
        }
        self.locked(|st, n| {
            if st.handle.is_some() || !st.listeners.is_empty() {
                return Err(Errno::EISCONN);
            }
            let local = if st.local_port != 0 { st.local_port } else { n.ephemeral_port() };
            let mut sock = new_tcp();
            sock.connect(n.iface.context(), endpoint(to), local).map_err(|_| Errno::EINVAL)?;
            st.handle = Some(n.sockets.add(sock));
            st.local_port = local;
            Ok(())
        })??;
        poll(); // send the SYN now
        let result = self.wait(Some(CONNECT_TIMEOUT_US), |st, n| {
            match n.sockets.get::<tcp::Socket>(st.handle?).state() {
                tcp::State::SynSent | tcp::State::SynReceived => None,
                tcp::State::Closed => Some(Err(Errno::ECONNREFUSED)), // got a RST
                _ => Some(Ok(())),
            }
        });
        if let Err(e) = result {
            // Failed: forget the half-open socket so connect can be retried.
            self.locked(|st, n| {
                if let Some(h) = st.handle.take() {
                    n.sockets.remove(h);
                }
            })?;
            return Err(if e == Errno::EAGAIN { Errno::ETIMEDOUT } else { e });
        }
        Ok(())
    }

    /// listen(): start `backlog` (1..=4) smoltcp listeners on our port.
    pub fn listen(&self, backlog: usize) -> Result<(), Errno> {
        if self.kind != Kind::Tcp {
            return Err(Errno::EOPNOTSUPP);
        }
        self.locked(|st, n| {
            if st.local_port == 0 || st.handle.is_some() {
                return Err(Errno::EINVAL);
            }
            while st.listeners.len() < backlog.clamp(1, MAX_BACKLOG) {
                let mut sock = new_tcp();
                sock.listen(st.local_port).map_err(|_| Errno::EINVAL)?;
                st.listeners.push(n.sockets.add(sock));
            }
            Ok(())
        })?
    }

    /// accept(): wait for one of our listeners to get a connection; it
    /// becomes the returned socket, a new listener takes its place.
    pub fn accept(&self) -> Result<(Socket, Addr), Errno> {
        if self.kind != Kind::Tcp {
            return Err(Errno::EOPNOTSUPP);
        }
        self.wait(self.recv_timeout(), |st, n| {
            if st.listeners.is_empty() {
                return Some(Err(Errno::EINVAL)); // not listening
            }
            let i = st.listeners.iter().position(|&h| {
                !matches!(n.sockets.get::<tcp::Socket>(h).state(), tcp::State::Listen | tcp::State::SynReceived)
            })?;
            let handle = st.listeners[i];
            let mut fresh = new_tcp();
            if fresh.listen(st.local_port).is_ok() {
                st.listeners[i] = n.sockets.add(fresh);
            } else {
                st.listeners.remove(i);
            }
            let peer = n.sockets.get::<tcp::Socket>(handle).remote_endpoint().map(addr).unwrap_or(([0; 4], 0));
            let state =
                State { handle: Some(handle), listeners: Vec::new(), local_port: st.local_port, remote: None, timeout_us: None };
            Some(Ok((Socket { kind: Kind::Tcp, state: Mutex::new(state) }, peer)))
        })
    }

    /// send()/sendto()/write(). Blocks while the send buffer is full.
    pub fn send(&self, data: &[u8], to: Option<Addr>) -> Result<usize, Errno> {
        if self.kind == Kind::Udp && without_interrupts(|| self.state.lock().local_port) == 0 {
            super::wait_configured(5_000_000)?;
            self.bind(0)?; // first send picks a local port, like Linux
        }
        let sent = self.wait(None, |st, n| {
            let Some(h) = st.handle else { return Some(Err(Errno::ENOTCONN)) };
            match self.kind {
                Kind::Tcp => {
                    let sock = n.sockets.get_mut::<tcp::Socket>(h);
                    if !sock.may_send() {
                        return Some(Err(Errno::EPIPE)); // closed
                    }
                    sock.can_send().then(|| sock.send_slice(data).map_err(|_| Errno::EPIPE))
                }
                Kind::Udp => {
                    let Some(dest) = to.map(endpoint).or(st.remote) else { return Some(Err(Errno::EDESTADDRREQ)) };
                    match n.sockets.get_mut::<udp::Socket>(h).send_slice(data, dest) {
                        Ok(()) => Some(Ok(data.len())),
                        Err(udp::SendError::BufferFull) => None,
                        Err(_) => Some(Err(Errno::EMSGSIZE)),
                    }
                }
                Kind::Icmp => {
                    let Some(dest) = to.map(endpoint).or(st.remote) else { return Some(Err(Errno::EDESTADDRREQ)) };
                    if data.len() < 8 {
                        return Some(Err(Errno::EINVAL)); // not even an ICMP header
                    }
                    // Our identifier goes in bytes 4..6 (smoltcp computes
                    // the checksum when it sends).
                    let mut packet = data.to_vec();
                    packet[4..6].copy_from_slice(&st.local_port.to_be_bytes());
                    match n.sockets.get_mut::<icmp::Socket>(h).send_slice(&packet, dest.addr) {
                        Ok(()) => Some(Ok(data.len())),
                        Err(icmp::SendError::BufferFull) => None,
                        Err(_) => Some(Err(Errno::EMSGSIZE)),
                    }
                }
            }
        })?;
        poll(); // put it on the wire now, not at netd's next wake-up
        Ok(sent)
    }

    /// recv()/recvfrom()/read(). Returns 0 at end of stream (TCP peer
    /// closed). A datagram bigger than `buf` is cut short (like Linux).
    pub fn recv(&self, buf: &mut [u8]) -> Result<(usize, Option<Addr>), Errno> {
        let result = self.wait(self.recv_timeout(), |st, n| {
            let Some(h) = st.handle else {
                return Some(Err(if self.kind == Kind::Tcp { Errno::ENOTCONN } else { Errno::EINVAL }));
            };
            match self.kind {
                Kind::Tcp => {
                    let sock = n.sockets.get_mut::<tcp::Socket>(h);
                    if sock.can_recv() {
                        Some(sock.recv_slice(buf).map(|len| (len, None)).map_err(|_| Errno::ECONNRESET))
                    } else if !sock.may_recv() {
                        Some(Ok((0, None))) // FIN received (or reset): end of stream
                    } else {
                        None
                    }
                }
                Kind::Udp => {
                    let sock = n.sockets.get_mut::<udp::Socket>(h);
                    if !sock.is_open() {
                        return Some(Err(Errno::EINVAL)); // never bound: nothing can arrive
                    }
                    let (data, meta) = sock.recv().ok()?;
                    let len = data.len().min(buf.len());
                    buf[..len].copy_from_slice(&data[..len]);
                    Some(Ok((len, Some(addr(meta.endpoint)))))
                }
                Kind::Icmp => {
                    let (data, from) = n.sockets.get_mut::<icmp::Socket>(h).recv().ok()?;
                    let len = data.len().min(buf.len());
                    buf[..len].copy_from_slice(&data[..len]);
                    Some(Ok((len, Some(addr(IpEndpoint::new(from, 0))))))
                }
            }
        })?;
        if self.kind == Kind::Tcp {
            poll(); // our receive window opened up: tell the peer
        }
        Ok(result)
    }

    /// Local and remote address (for `accept` callers and tests).
    pub fn peer(&self) -> Option<Addr> {
        self.locked(|st, n| match (self.kind, st.handle) {
            (Kind::Tcp, Some(h)) => n.sockets.get::<tcp::Socket>(h).remote_endpoint().map(addr),
            _ => st.remote.map(addr),
        })
        .ok()
        .flatten()
    }
}

impl Drop for Socket {
    /// close(): TCP sends its FIN and lingers (net::reap_closing frees it
    /// once the goodbye handshake is done); the rest is freed at once.
    fn drop(&mut self) {
        let kind = self.kind;
        let _ = self.locked(|st, n| {
            for h in st.listeners.drain(..) {
                n.sockets.remove(h);
            }
            if let Some(h) = st.handle.take() {
                if kind == Kind::Tcp {
                    let sock = n.sockets.get_mut::<tcp::Socket>(h);
                    sock.close();
                    sock.set_timeout(Some(smoltcp::time::Duration::from_millis(CLOSE_TIMEOUT_MS)));
                    n.close_later(h);
                } else {
                    n.sockets.remove(h);
                }
            }
        });
        super::kick(); // let netd send the FIN (we may hold other locks here)
    }
}

impl File for Socket {
    fn read(&self, buf: &mut [u8]) -> Result<usize, Errno> {
        self.recv(buf).map(|(n, _)| n)
    }
    fn write(&self, buf: &[u8]) -> Result<usize, Errno> {
        self.send(buf, None)
    }
    fn stat(&self) -> Result<Stat, Errno> {
        Ok(Stat { kind: KIND_SOCKET, size: 0, mode: 0 })
    }
    fn as_socket(&self) -> Option<&Socket> {
        Some(self)
    }
}