3649 lines · 127.3 KB · Rust
RawDownload
use crate::fd_table::TransferredFd;use crate::poll::{PollEvents, POLLERR, POLLHUP, POLLIN, POLLOUT};use crate::vfs::normalize_path;#[cfg(not(target_arch = "wasm32"))]use agentos_runtime::accounting::{Reservation, ResourceClass, ResourceLedger};use std::any::Any;use std::collections::{BTreeMap, BTreeSet, VecDeque};use std::error::Error;use std::fmt;use std::net::{Ipv4Addr, Ipv6Addr};use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};use std::sync::{Arc, Mutex, MutexGuard};use std::time::Instant; pub type SocketId = u64;pub type SocketResult<T> = Result<T, SocketTableError>; #[derive(Debug, Clone, Copy, PartialEq, Eq)]pub enum SocketReadinessKind {    Data,    Accept,} #[derive(Debug, Clone, Copy, PartialEq, Eq)]pub struct SocketReadiness {    pub socket_id: SocketId,    pub kind: SocketReadinessKind,} type SocketReadinessSink = Arc<dyn Fn(SocketReadiness) + Send + Sync + 'static>; #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]pub struct SocketReadTraceSnapshot {    pub socket_record_clone_calls: u64,    pub socket_record_clone_us: u64,    pub read_recv_calls: u64,    pub read_recv_bytes: u64,    pub read_recv_chunks: u64,    pub read_recv_copy_us: u64,} struct SocketReadTraceCounters {    socket_record_clone_calls: AtomicU64,    socket_record_clone_us: AtomicU64,    read_recv_calls: AtomicU64,    read_recv_bytes: AtomicU64,    read_recv_chunks: AtomicU64,    read_recv_copy_us: AtomicU64,} impl SocketReadTraceCounters {    const fn new() -> Self {        Self {            socket_record_clone_calls: AtomicU64::new(0),            socket_record_clone_us: AtomicU64::new(0),            read_recv_calls: AtomicU64::new(0),            read_recv_bytes: AtomicU64::new(0),            read_recv_chunks: AtomicU64::new(0),            read_recv_copy_us: AtomicU64::new(0),        }    }} static SOCKET_READ_TRACE_ENABLED: AtomicBool = AtomicBool::new(false);static SOCKET_READ_TRACE_COUNTERS: SocketReadTraceCounters = SocketReadTraceCounters::new(); pub fn set_socket_read_trace_enabled(enabled: bool) {    SOCKET_READ_TRACE_ENABLED.store(enabled, Ordering::Relaxed);} pub fn reset_socket_read_trace() {    for counter in [        &SOCKET_READ_TRACE_COUNTERS.socket_record_clone_calls,        &SOCKET_READ_TRACE_COUNTERS.socket_record_clone_us,        &SOCKET_READ_TRACE_COUNTERS.read_recv_calls,        &SOCKET_READ_TRACE_COUNTERS.read_recv_bytes,        &SOCKET_READ_TRACE_COUNTERS.read_recv_chunks,        &SOCKET_READ_TRACE_COUNTERS.read_recv_copy_us,    ] {        counter.store(0, Ordering::Relaxed);    }} pub fn socket_read_trace_snapshot() -> SocketReadTraceSnapshot {    SocketReadTraceSnapshot {        socket_record_clone_calls: SOCKET_READ_TRACE_COUNTERS            .socket_record_clone_calls            .load(Ordering::Relaxed),        socket_record_clone_us: SOCKET_READ_TRACE_COUNTERS            .socket_record_clone_us            .load(Ordering::Relaxed),        read_recv_calls: SOCKET_READ_TRACE_COUNTERS            .read_recv_calls            .load(Ordering::Relaxed),        read_recv_bytes: SOCKET_READ_TRACE_COUNTERS            .read_recv_bytes            .load(Ordering::Relaxed),        read_recv_chunks: SOCKET_READ_TRACE_COUNTERS            .read_recv_chunks            .load(Ordering::Relaxed),        read_recv_copy_us: SOCKET_READ_TRACE_COUNTERS            .read_recv_copy_us            .load(Ordering::Relaxed),    }} #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]pub struct InetSocketAddress {    host: String,    port: u16,} impl InetSocketAddress {    pub fn new(host: impl Into<String>, port: u16) -> Self {        Self {            host: host.into(),            port,        }    }     pub fn host(&self) -> &str {        &self.host    }     pub const fn port(&self) -> u16 {        self.port    }} #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]pub enum SocketDomain {    Inet,    Inet6,    Unix,} #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]pub enum SocketType {    Stream,    Datagram,    SeqPacket,} #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]pub enum SocketState {    Created,    Bound,    Listening,    Connected,} impl SocketState {    pub const fn counts_as_listener(self) -> bool {        matches!(self, Self::Listening)    }     pub const fn counts_as_connection(self) -> bool {        matches!(self, Self::Connected)    }} #[derive(Debug, Clone, Copy, PartialEq, Eq)]pub enum SocketShutdown {    Read,    Write,    Both,} #[derive(Debug, Clone, Copy, PartialEq, Eq)]pub enum DatagramSocketOption {    ReuseAddr,    ReusePort,    Broadcast,} #[derive(Debug, Clone, Copy, PartialEq, Eq)]pub struct SocketSpec {    pub domain: SocketDomain,    pub socket_type: SocketType,} impl SocketSpec {    pub const fn new(domain: SocketDomain, socket_type: SocketType) -> Self {        Self {            domain,            socket_type,        }    }     pub const fn tcp() -> Self {        Self::new(SocketDomain::Inet, SocketType::Stream)    }     pub const fn udp() -> Self {        Self::new(SocketDomain::Inet, SocketType::Datagram)    }     pub const fn unix_stream() -> Self {        Self::new(SocketDomain::Unix, SocketType::Stream)    }     pub const fn unix_datagram() -> Self {        Self::new(SocketDomain::Unix, SocketType::Datagram)    }     pub const fn unix_seqpacket() -> Self {        Self::new(SocketDomain::Unix, SocketType::SeqPacket)    }} #[derive(Debug, Clone, PartialEq, Eq)]pub struct SocketRecord {    id: SocketId,    owner_pid: u32,    spec: SocketSpec,    state: SocketState,    local_address: Option<InetSocketAddress>,    peer_address: Option<InetSocketAddress>,    local_unix_path: Option<String>,    peer_unix_path: Option<String>,    listener_state: Option<ListenerState>,    connection_state: Option<ConnectionState>,    datagram_state: Option<DatagramState>,} impl SocketRecord {    pub const fn id(&self) -> SocketId {        self.id    }     pub const fn owner_pid(&self) -> u32 {        self.owner_pid    }     pub const fn spec(&self) -> SocketSpec {        self.spec    }     pub const fn state(&self) -> SocketState {        self.state    }     pub fn local_address(&self) -> Option<&InetSocketAddress> {        self.local_address.as_ref()    }     pub fn peer_address(&self) -> Option<&InetSocketAddress> {        self.peer_address.as_ref()    }     pub fn local_unix_path(&self) -> Option<&str> {        self.local_unix_path.as_deref()    }     pub fn peer_unix_path(&self) -> Option<&str> {        self.peer_unix_path.as_deref()    }     pub fn listen_backlog(&self) -> Option<usize> {        self.listener_state.as_ref().map(|state| state.backlog)    }     pub fn pending_accept_count(&self) -> usize {        self.listener_state            .as_ref()            .map(|state| state.pending_accepts.len())            .unwrap_or(0)    }     pub fn peer_socket_id(&self) -> Option<SocketId> {        self.connection_state            .as_ref()            .and_then(|state| state.peer_socket_id)    }     pub fn buffered_read_bytes(&self) -> usize {        self.connection_state            .as_ref()            .map(ConnectionState::buffered_len)            .unwrap_or(0)    }     pub fn read_shutdown(&self) -> bool {        self.connection_state            .as_ref()            .map(|state| state.read_shutdown)            .unwrap_or(false)    }     pub fn write_shutdown(&self) -> bool {        self.connection_state            .as_ref()            .map(|state| state.write_shutdown)            .unwrap_or(false)    }     pub fn peer_write_shutdown(&self) -> bool {        self.connection_state            .as_ref()            .map(|state| state.peer_write_shutdown)            .unwrap_or(false)    }     pub fn queued_datagrams(&self) -> usize {        self.datagram_state            .as_ref()            .map(|state| state.recv_queue.len())            .unwrap_or_else(|| {                self.connection_state                    .as_ref()                    .filter(|_| self.spec.socket_type != SocketType::Stream)                    .map(|state| state.recv_buffer.len())                    .unwrap_or(0)            })    }     pub fn queued_datagram_bytes(&self) -> usize {        self.datagram_state            .as_ref()            .map(|state| datagram_queue_bytes(&state.recv_queue))            .unwrap_or(0)    }     pub fn reuse_address(&self) -> bool {        self.datagram_state            .as_ref()            .map(|state| state.reuse_addr)            .unwrap_or(false)    }     pub fn reuse_port(&self) -> bool {        self.datagram_state            .as_ref()            .map(|state| state.reuse_port)            .unwrap_or(false)    }     pub fn broadcast_enabled(&self) -> bool {        self.datagram_state            .as_ref()            .map(|state| state.broadcast)            .unwrap_or(false)    }     pub fn multicast_membership_count(&self) -> usize {        self.datagram_state            .as_ref()            .map(|state| state.multicast_memberships.len())            .unwrap_or(0)    }     pub fn has_multicast_membership(&self, membership: &SocketMulticastMembership) -> bool {        self.datagram_state            .as_ref()            .map(|state| state.multicast_memberships.contains(membership))            .unwrap_or(false)    }} #[derive(Debug, Clone, PartialEq, Eq)]pub struct ReceivedDatagram {    source_address: Option<InetSocketAddress>,    payload: Vec<u8>,} pub type OpaqueTransferredRight = Arc<dyn Any + Send + Sync + 'static>; #[derive(Clone)]pub enum TransferredSocketRight {    Fd(TransferredFd),    Opaque(OpaqueTransferredRight),} impl fmt::Debug for TransferredSocketRight {    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {        match self {            Self::Fd(fd) => f.debug_tuple("Fd").field(&fd.description_id()).finish(),            Self::Opaque(resource) => f                .debug_tuple("Opaque")                .field(&(Arc::as_ptr(resource) as *const ()))                .finish(),        }    }} impl PartialEq for TransferredSocketRight {    fn eq(&self, other: &Self) -> bool {        match (self, other) {            (Self::Fd(left), Self::Fd(right)) => left == right,            (Self::Opaque(left), Self::Opaque(right)) => Arc::ptr_eq(left, right),            _ => false,        }    }} impl Eq for TransferredSocketRight {} #[derive(Debug)]pub struct ReceivedSocketMessage {    pub payload: Vec<u8>,    pub rights: Vec<TransferredSocketRight>,    pub truncated: bool,    pub full_length: usize,} impl ReceivedDatagram {    pub fn source_address(&self) -> Option<&InetSocketAddress> {        self.source_address.as_ref()    }     pub fn payload(&self) -> &[u8] {        &self.payload    }     pub fn into_parts(self) -> (Option<InetSocketAddress>, Vec<u8>) {        (self.source_address, self.payload)    }} /// Native sidecar handoff that transfers the kernel queue's exact resource/// ownership with the datagram. Standalone callers may keep using/// `recv_datagram`, whose boundary releases queue ownership immediately.#[cfg(not(target_arch = "wasm32"))]pub type DatagramReservations = (Reservation, Reservation, Reservation, Reservation); #[cfg(not(target_arch = "wasm32"))]#[derive(Debug)]pub struct ChargedReceivedDatagram {    datagram: ReceivedDatagram,    reservations: Option<DatagramReservations>,} #[cfg(not(target_arch = "wasm32"))]impl ChargedReceivedDatagram {    pub fn into_parts(        self,    ) -> (        Option<InetSocketAddress>,        Vec<u8>,        Option<DatagramReservations>,    ) {        let (source_address, payload) = self.datagram.into_parts();        (source_address, payload, self.reservations)    }} #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]pub struct SocketTableSnapshot {    pub sockets: usize,    pub listeners: usize,    pub connections: usize,    pub buffered_bytes: usize,    pub datagram_queue_len: usize,} #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]pub struct SocketMulticastMembership {    group_address: String,    interface_address: Option<String>,} impl SocketMulticastMembership {    pub fn new(group_address: impl Into<String>, interface_address: Option<String>) -> Self {        Self {            group_address: group_address.into(),            interface_address,        }    }     pub fn group_address(&self) -> &str {        &self.group_address    }     pub fn interface_address(&self) -> Option<&str> {        self.interface_address.as_deref()    }} #[derive(Debug, Clone, PartialEq, Eq)]pub struct SocketTableError {    code: &'static str,    message: String,} impl SocketTableError {    pub fn code(&self) -> &'static str {        self.code    }     fn not_found(socket_id: SocketId) -> Self {        Self {            code: "ENOENT",            message: format!("no such socket {socket_id}"),        }    }     fn invalid_argument(message: impl Into<String>) -> Self {        Self {            code: "EINVAL",            message: message.into(),

Lines 1–500 of 3649.