Skip to main content

ax_net/
socket.rs

1//! Public socket facade.
2//!
3//! This module defines the protocol-independent socket API used by syscall
4//! layers: common send/recv flags, extended address families, shutdown modes,
5//! and the `SocketOps` trait implemented by TCP, UDP, raw, Unix, and vsock
6//! transports.
7//!
8//! # Compatibility Boundary
9//!
10//! The syscall layer should not need to know whether a socket is backed by
11//! smoltcp, an in-kernel Unix transport, or a vsock connection manager. It
12//! passes `SocketAddrEx`, `SendOptions`, and `RecvOptions` into this facade, and
13//! each concrete transport maps them onto its own semantics.
14//!
15//! # Design Rule
16//!
17//! This module contains dispatch and common ABI shapes only. Protocol behavior
18//! such as TCP accept queues, UDP corking, raw packet format, or Unix ancillary
19//! data delivery belongs in the corresponding transport module.
20
21use alloc::{boxed::Box, vec::Vec};
22use core::{
23    any::Any,
24    fmt::{self, Debug},
25    net::SocketAddr,
26    time::Duration,
27};
28
29use ax_io::prelude::*;
30use axpoll::{ExclusiveRegistrationSink, IoEvents, Pollable, SharedRegistrationSink};
31use bitflags::bitflags;
32use enum_dispatch::enum_dispatch;
33
34#[cfg(feature = "vsock")]
35use crate::vsock::{VsockAddr, VsockSocket};
36use crate::{
37    NetError, NetResult,
38    options::{Configurable, GetSocketOption, SetSocketOption, UnixCredentials},
39    raw::RawSocket,
40    tcp::TcpSocket,
41    udp::UdpSocket,
42    unix::{UnixSocket, UnixSocketAddr},
43};
44
45/// Extended socket address supporting IP, Unix, and vsock address families.
46#[derive(Clone, Debug)]
47pub enum SocketAddrEx {
48    /// An IP (v4/v6) socket address.
49    Ip(SocketAddr),
50    /// A Unix domain socket address.
51    Unix(UnixSocketAddr),
52    /// A vsock socket address.
53    #[cfg(feature = "vsock")]
54    Vsock(VsockAddr),
55}
56
57impl SocketAddrEx {
58    /// Convert into an IP socket address, or return an error if not IP.
59    pub fn into_ip(self) -> NetResult<SocketAddr> {
60        match self {
61            SocketAddrEx::Ip(addr) => Ok(addr),
62            SocketAddrEx::Unix(_) => Err(NetError::AddressFamilyUnsupported),
63            #[cfg(feature = "vsock")]
64            SocketAddrEx::Vsock(_) => Err(NetError::AddressFamilyUnsupported),
65        }
66    }
67
68    /// Convert into a Unix socket address, or return an error if not Unix.
69    pub fn into_unix(self) -> NetResult<UnixSocketAddr> {
70        match self {
71            SocketAddrEx::Unix(addr) => Ok(addr),
72            SocketAddrEx::Ip(_) => Err(NetError::AddressFamilyUnsupported),
73            #[cfg(feature = "vsock")]
74            SocketAddrEx::Vsock(_) => Err(NetError::AddressFamilyUnsupported),
75        }
76    }
77
78    /// Convert into a vsock address, or return an error if not vsock.
79    #[cfg(feature = "vsock")]
80    pub fn into_vsock(self) -> NetResult<VsockAddr> {
81        match self {
82            SocketAddrEx::Ip(_) => Err(NetError::AddressFamilyUnsupported),
83            SocketAddrEx::Unix(_) => Err(NetError::AddressFamilyUnsupported),
84            SocketAddrEx::Vsock(addr) => Ok(addr),
85        }
86    }
87}
88
89bitflags! {
90    /// Flags for sending data to a socket.
91    ///
92    /// These values match Linux MSG_* constants so that `from_bits_retain(flags)`
93    /// from the syscall layer preserves the correct flags.
94    ///
95    /// See [`SocketOps::send`].
96    #[derive(Default, Debug, Clone, Copy)]
97    pub struct SendFlags: u32 {
98        /// Sends out-of-band data on sockets that support it (e.g. SOCK_STREAM).
99        const OOB = 0x01;
100        /// Don't use a gateway to send the packet, send to hosts only on
101        /// directly connected networks.
102        const DONTROUTE = 0x04;
103        /// Enables nonblocking operation; if the operation would block,
104        /// `EAGAIN` or `EWOULDBLOCK` is returned.
105        const DONTWAIT = 0x40;
106        /// Terminates a record (SOCK_SEQPACKET).
107        const EOR = 0x80;
108        /// Sends only if a connection confirm is pending (UDP/RAW, Linux specific).
109        const CONFIRM = 0x800;
110        /// Requests not to send SIGPIPE on errors on stream oriented sockets
111        /// when the other end breaks the connection.
112        const NOSIGNAL = 0x4000;
113        /// More data will be sent; used to cork/coalesce sends (UDP/TCP).
114        const MORE = 0x8000;
115    }
116}
117
118bitflags! {
119    /// Flags for receiving data from a socket.
120    ///
121    /// See [`SocketOps::recv`].
122    #[derive(Default, Debug, Clone, Copy)]
123    pub struct RecvFlags: u32 {
124        /// Receive data without removing it from the queue.
125        const PEEK = 0x01;
126        /// For datagram-like sockets, requires [`SocketOps::recv`] to return
127        /// the real size of the datagram, even when it is larger than the
128        /// buffer.
129        const TRUNCATE = 0x02;
130        /// Requests out-of-band data (`MSG_OOB`). Only stream sockets that
131        /// support urgent data honor it; others reject it with `EOPNOTSUPP`.
132        const OOB = 0x04;
133        /// Per-call non-blocking override (`MSG_DONTWAIT`). Does NOT
134        /// change the socket's own `O_NONBLOCK` state.
135        const DONTWAIT = 0x40;
136    }
137}
138
139/// Ancillary control message payload, carried opaquely so the socket layer
140/// stays protocol-independent. Cloneable so `recvmsg(MSG_PEEK)` can duplicate
141/// the ancillary data without consuming the record: SCM_RIGHTS fds are cloned
142/// (sharing the open file description), matching Linux `unix_peek_fds` /
143/// `scm_fp_dup`.
144pub trait CMsgPayload: Any + Send + Sync {
145    /// Duplicate into a fresh payload for peek delivery.
146    fn clone_box(&self) -> Box<dyn CMsgPayload>;
147    /// Recover a `Box<dyn Any>` for owned downcast at the syscall layer.
148    fn into_any(self: Box<Self>) -> Box<dyn Any + Send + Sync>;
149}
150impl<T: Any + Send + Sync + Clone> CMsgPayload for T {
151    fn clone_box(&self) -> Box<dyn CMsgPayload> {
152        Box::new(self.clone())
153    }
154    fn into_any(self: Box<Self>) -> Box<dyn Any + Send + Sync> {
155        self
156    }
157}
158// Opaque payload; mirror the std `dyn Any` Debug impl so containers deriving
159// Debug still compile.
160impl core::fmt::Debug for dyn CMsgPayload {
161    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
162        f.write_str("CMsgPayload { .. }")
163    }
164}
165
166/// Type alias for ancillary control message data.
167pub type CMsgData = Box<dyn CMsgPayload>;
168
169/// IP ancillary data reported through `recvmsg`.
170#[derive(Debug, Clone, Copy, PartialEq, Eq)]
171pub enum IpCmsg {
172    /// IPv4 hop limit for `IP_RECVTTL`.
173    Ipv4Ttl(u8),
174    /// IPv4 TOS byte for `IP_RECVTOS`.
175    Ipv4Tos(u8),
176    /// IPv6 traffic-class byte for `IPV6_RECVTCLASS`.
177    Ipv6TrafficClass(u8),
178}
179
180/// Transport-independent socket-level ancillary data reported through
181/// `recvmsg`.
182#[derive(Debug, Clone, PartialEq, Eq)]
183pub enum SocketCmsg {
184    /// Sender credentials requested with `SO_PASSCRED`.
185    Credentials(UnixCredentials),
186    /// Wall-clock receive timestamp requested with `SO_TIMESTAMP`.
187    Timestamp(Duration),
188}
189
190/// Options for sending data to a socket.
191///
192/// See [`SocketOps::send`].
193#[derive(Default, Debug)]
194pub struct SendOptions {
195    /// Destination address for the message.
196    pub to: Option<SocketAddrEx>,
197    /// Send flags.
198    pub flags: SendFlags,
199    /// Ancillary control messages.
200    pub cmsg: Vec<CMsgData>,
201    /// Real credentials of the task performing this send operation.
202    pub sender_credentials: Option<UnixCredentials>,
203}
204
205/// Options for receiving data from a socket.
206///
207/// See [`SocketOps::recv`].
208#[derive(Default)]
209pub struct RecvOptions<'a> {
210    /// If set, the sender's address is written here.
211    pub from: Option<&'a mut SocketAddrEx>,
212    /// Receive flags.
213    pub flags: RecvFlags,
214    /// If set, ancillary control messages are appended here.
215    pub cmsg: Option<&'a mut Vec<CMsgData>>,
216    /// If set and the datagram was truncated, this is set to `true`.
217    pub truncated: Option<&'a mut bool>,
218}
219impl Debug for RecvOptions<'_> {
220    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
221        f.debug_struct("RecvOptions")
222            .field("from", &self.from)
223            .field("flags", &self.flags)
224            .finish()
225    }
226}
227
228/// Kind of shutdown operation to perform on a socket.
229#[derive(Debug, Clone, Copy)]
230pub enum Shutdown {
231    /// Shut down the read half.
232    Read,
233    /// Shut down the write half.
234    Write,
235    /// Shut down both halves.
236    Both,
237}
238
239/// Progress of a connection attempt whose completion is readiness-driven.
240#[derive(Clone, Copy, Debug, Eq, PartialEq)]
241pub enum ConnectStatus {
242    /// The transport completed the connection synchronously or asynchronously.
243    Connected,
244    /// The transport started an asynchronous connection and is not ready yet.
245    InProgress,
246}
247
248/// OS-facing policy for driving one task-neutral socket future.
249#[derive(Clone, Copy, Debug, Eq, PartialEq)]
250pub struct SocketWaitPolicy {
251    /// Whether this operation must return after the registration recheck.
252    pub nonblocking: bool,
253    /// Optional relative timeout owned and enforced by the consuming OS.
254    pub timeout: Option<Duration>,
255}
256
257fn socket_wait_policy(
258    socket: &(impl Configurable + ?Sized),
259    send: bool,
260    extra_nonblocking: bool,
261) -> NetResult<SocketWaitPolicy> {
262    let mut nonblocking = false;
263    socket.get_option(GetSocketOption::NonBlocking(&mut nonblocking))?;
264    let mut timeout = Duration::ZERO;
265    if send {
266        socket.get_option(GetSocketOption::SendTimeout(&mut timeout))?;
267    } else {
268        socket.get_option(GetSocketOption::ReceiveTimeout(&mut timeout))?;
269    }
270    Ok(SocketWaitPolicy {
271        nonblocking: nonblocking || extra_nonblocking,
272        timeout: (!timeout.is_zero()).then_some(timeout),
273    })
274}
275impl Shutdown {
276    /// Returns `true` if the read half should be shut down.
277    pub fn has_read(&self) -> bool {
278        matches!(self, Shutdown::Read | Shutdown::Both)
279    }
280
281    /// Returns `true` if the write half should be shut down.
282    pub fn has_write(&self) -> bool {
283        matches!(self, Shutdown::Write | Shutdown::Both)
284    }
285}
286
287/// Operations that can be performed on a socket.
288#[enum_dispatch]
289pub trait SocketOps: Configurable {
290    /// Binds an unbound socket to the given address and port.
291    fn bind(&self, local_addr: SocketAddrEx) -> NetResult;
292    /// Starts a connection without waiting for asynchronous completion.
293    fn start_connect(&self, remote_addr: SocketAddrEx) -> NetResult<ConnectStatus>;
294    /// Checks a previously started asynchronous connection.
295    fn connect_status(&self) -> NetResult<ConnectStatus> {
296        Ok(ConnectStatus::Connected)
297    }
298
299    /// Starts listening on the bound address and port.
300    fn listen(&self, _backlog: usize) -> NetResult {
301        Err(NetError::OperationNotSupported)
302    }
303    /// Returns whether this socket currently accepts incoming connections.
304    fn is_listening(&self) -> bool {
305        false
306    }
307    /// Attempts to accept one connection without parking the caller.
308    fn try_accept(&self) -> NetResult<Socket> {
309        Err(NetError::OperationNotSupported)
310    }
311
312    /// Attempts to send data without parking the caller.
313    fn try_send(&self, src: impl Read + IoBuf, options: &mut SendOptions) -> NetResult<usize>;
314    /// Attempts to receive data without parking the caller.
315    fn try_recv(
316        &self,
317        dst: impl Write + IoBufMut,
318        options: &mut RecvOptions<'_>,
319    ) -> NetResult<usize>;
320    /// Returns the number of bytes that can be read without blocking.
321    fn recv_available(&self) -> NetResult<usize> {
322        Err(NetError::OperationNotSupported)
323    }
324
325    /// Get the local endpoint of the socket.
326    fn local_addr(&self) -> NetResult<SocketAddrEx>;
327    /// Get the remote endpoint of the socket.
328    fn peer_addr(&self) -> NetResult<SocketAddrEx>;
329
330    /// Shutdown the socket, closing the connection.
331    fn shutdown(&self, how: Shutdown) -> NetResult;
332
333    /// Returns the send wait policy without acquiring scheduler ownership.
334    fn send_wait_policy(&self, extra_nonblocking: bool) -> NetResult<SocketWaitPolicy> {
335        socket_wait_policy(self, true, extra_nonblocking)
336    }
337
338    /// Returns the receive wait policy without acquiring scheduler ownership.
339    fn recv_wait_policy(&self, extra_nonblocking: bool) -> NetResult<SocketWaitPolicy> {
340        socket_wait_policy(self, false, extra_nonblocking)
341    }
342}
343
344impl<T: SocketOps + ?Sized> SocketOps for Box<T> {
345    fn bind(&self, local_addr: SocketAddrEx) -> NetResult {
346        (**self).bind(local_addr)
347    }
348
349    fn start_connect(&self, remote_addr: SocketAddrEx) -> NetResult<ConnectStatus> {
350        (**self).start_connect(remote_addr)
351    }
352
353    fn connect_status(&self) -> NetResult<ConnectStatus> {
354        (**self).connect_status()
355    }
356
357    fn listen(&self, backlog: usize) -> NetResult {
358        (**self).listen(backlog)
359    }
360
361    fn is_listening(&self) -> bool {
362        (**self).is_listening()
363    }
364
365    fn try_accept(&self) -> NetResult<Socket> {
366        (**self).try_accept()
367    }
368
369    fn try_send(&self, src: impl Read + IoBuf, options: &mut SendOptions) -> NetResult<usize> {
370        (**self).try_send(src, options)
371    }
372
373    fn try_recv(
374        &self,
375        dst: impl Write + IoBufMut,
376        options: &mut RecvOptions<'_>,
377    ) -> NetResult<usize> {
378        (**self).try_recv(dst, options)
379    }
380
381    fn recv_available(&self) -> NetResult<usize> {
382        (**self).recv_available()
383    }
384
385    fn local_addr(&self) -> NetResult<SocketAddrEx> {
386        (**self).local_addr()
387    }
388
389    fn peer_addr(&self) -> NetResult<SocketAddrEx> {
390        (**self).peer_addr()
391    }
392
393    fn shutdown(&self, how: Shutdown) -> NetResult {
394        (**self).shutdown(how)
395    }
396}
397
398/// Network socket abstraction.
399#[enum_dispatch(Configurable, SocketOps)]
400pub enum Socket {
401    /// UDP socket.
402    Udp(Box<UdpSocket>),
403    /// TCP socket.
404    Tcp(Box<TcpSocket>),
405    /// Raw IP socket.
406    Raw(Box<RawSocket>),
407    /// Unix domain socket.
408    Unix(Box<UnixSocket>),
409    /// Virtio socket.
410    #[cfg(feature = "vsock")]
411    Vsock(Box<VsockSocket>),
412}
413
414impl From<UdpSocket> for Socket {
415    fn from(socket: UdpSocket) -> Self {
416        Self::Udp(Box::new(socket))
417    }
418}
419
420impl From<TcpSocket> for Socket {
421    fn from(socket: TcpSocket) -> Self {
422        Self::Tcp(Box::new(socket))
423    }
424}
425
426impl From<UnixSocket> for Socket {
427    fn from(socket: UnixSocket) -> Self {
428        Self::Unix(Box::new(socket))
429    }
430}
431
432#[cfg(feature = "vsock")]
433impl From<VsockSocket> for Socket {
434    fn from(socket: VsockSocket) -> Self {
435        Self::Vsock(Box::new(socket))
436    }
437}
438
439impl Pollable for Socket {
440    fn poll(&self) -> IoEvents {
441        match self {
442            Socket::Tcp(tcp) => tcp.poll(),
443            Socket::Udp(udp) => udp.poll(),
444            Socket::Raw(raw) => raw.poll(),
445            Socket::Unix(unix) => unix.poll(),
446            #[cfg(feature = "vsock")]
447            Socket::Vsock(vsock) => vsock.poll(),
448        }
449    }
450
451    unsafe fn register_shared(&self, sink: &mut dyn SharedRegistrationSink, events: IoEvents) {
452        match self {
453            Socket::Tcp(tcp) => unsafe { tcp.register_shared(sink, events) },
454            Socket::Udp(udp) => unsafe { udp.register_shared(sink, events) },
455            Socket::Raw(raw) => unsafe { raw.register_shared(sink, events) },
456            Socket::Unix(unix) => unsafe { unix.register_shared(sink, events) },
457            #[cfg(feature = "vsock")]
458            Socket::Vsock(vsock) => unsafe { vsock.register_shared(sink, events) },
459        }
460    }
461
462    unsafe fn register_exclusive(
463        &self,
464        sink: &mut dyn ExclusiveRegistrationSink,
465        events: IoEvents,
466    ) {
467        match self {
468            Socket::Tcp(tcp) => unsafe { tcp.register_exclusive(sink, events) },
469            Socket::Udp(udp) => unsafe { udp.register_exclusive(sink, events) },
470            Socket::Raw(raw) => unsafe { raw.register_exclusive(sink, events) },
471            Socket::Unix(unix) => unsafe { unix.register_exclusive(sink, events) },
472            #[cfg(feature = "vsock")]
473            Socket::Vsock(vsock) => unsafe { vsock.register_exclusive(sink, events) },
474        }
475    }
476}