os/net/
socket.rs

1//! Socket implementation using smoltcp
2
3use crate::sync::SpinLock;
4use crate::vfs::{File, FsError, InodeMetadata};
5use alloc::vec;
6use lazy_static::lazy_static;
7use smoltcp::iface::{Interface, SocketHandle as SmoltcpHandle, SocketSet};
8use smoltcp::socket::{tcp, udp};
9use smoltcp::wire::{IpAddress, IpEndpoint, Ipv4Address};
10
11#[derive(Clone, Copy)]
12pub enum SocketHandle {
13    Tcp(SmoltcpHandle),
14    Udp(SmoltcpHandle),
15}
16
17use alloc::collections::BTreeMap;
18use alloc::sync::Arc;
19
20lazy_static! {
21    pub static ref SOCKET_SET: SpinLock<SocketSet<'static>> = SpinLock::new(SocketSet::new(vec![]));
22    /// Map from fd to socket handle for syscall operations
23    pub static ref FD_SOCKET_MAP: SpinLock<BTreeMap<(usize, usize), SocketHandle>> = SpinLock::new(BTreeMap::new());
24    /// Global network interface for socket operations
25    pub static ref GLOBAL_INTERFACE: SpinLock<Option<Interface>> = SpinLock::new(None);
26}
27
28use crate::uapi::fcntl::OpenFlags;
29use crate::uapi::socket::SocketOptions;
30
31pub struct SocketFile {
32    handle: SpinLock<SocketHandle>,
33    local_endpoint: SpinLock<Option<IpEndpoint>>,
34    remote_endpoint: SpinLock<Option<IpEndpoint>>,
35    shutdown_rd: SpinLock<bool>,
36    shutdown_wr: SpinLock<bool>,
37    flags: SpinLock<OpenFlags>,
38    options: SpinLock<SocketOptions>,
39}
40
41impl SocketFile {
42    pub fn new(handle: SocketHandle) -> Self {
43        Self {
44            handle: SpinLock::new(handle),
45            local_endpoint: SpinLock::new(None),
46            remote_endpoint: SpinLock::new(None),
47            shutdown_rd: SpinLock::new(false),
48            shutdown_wr: SpinLock::new(false),
49            flags: SpinLock::new(OpenFlags::empty()),
50            options: SpinLock::new(SocketOptions::default()),
51        }
52    }
53
54    pub fn new_with_flags(handle: SocketHandle, flags: OpenFlags) -> Self {
55        Self {
56            handle: SpinLock::new(handle),
57            local_endpoint: SpinLock::new(None),
58            remote_endpoint: SpinLock::new(None),
59            shutdown_rd: SpinLock::new(false),
60            shutdown_wr: SpinLock::new(false),
61            flags: SpinLock::new(flags),
62            options: SpinLock::new(SocketOptions::default()),
63        }
64    }
65
66    pub fn get_socket_options(&self) -> SocketOptions {
67        *self.options.lock()
68    }
69
70    pub fn set_socket_options(&self, opts: SocketOptions) {
71        *self.options.lock() = opts;
72    }
73
74    pub fn handle(&self) -> SocketHandle {
75        *self.handle.lock()
76    }
77
78    pub fn set_handle(&self, new_handle: SocketHandle) {
79        *self.handle.lock() = new_handle;
80    }
81
82    pub fn set_local_endpoint(&self, endpoint: IpEndpoint) {
83        *self.local_endpoint.lock() = Some(endpoint);
84    }
85
86    pub fn get_local_endpoint(&self) -> Option<IpEndpoint> {
87        *self.local_endpoint.lock()
88    }
89
90    pub fn set_remote_endpoint(&self, endpoint: IpEndpoint) {
91        *self.remote_endpoint.lock() = Some(endpoint);
92    }
93
94    pub fn get_remote_endpoint(&self) -> Option<IpEndpoint> {
95        *self.remote_endpoint.lock()
96    }
97
98    pub fn shutdown_read(&self) {
99        *self.shutdown_rd.lock() = true;
100    }
101
102    pub fn shutdown_write(&self) {
103        *self.shutdown_wr.lock() = true;
104    }
105
106    pub fn is_shutdown_read(&self) -> bool {
107        *self.shutdown_rd.lock()
108    }
109
110    pub fn is_shutdown_write(&self) -> bool {
111        *self.shutdown_wr.lock()
112    }
113}
114
115impl Drop for SocketFile {
116    fn drop(&mut self) {
117        let mut sockets = SOCKET_SET.lock();
118        match *self.handle.lock() {
119            SocketHandle::Tcp(h) => {
120                sockets.remove(h);
121            }
122            SocketHandle::Udp(h) => {
123                sockets.remove(h);
124            }
125        }
126    }
127}
128
129/// Register a socket fd mapping (tid, fd) -> handle
130pub fn register_socket_fd(tid: usize, fd: usize, handle: SocketHandle) {
131    FD_SOCKET_MAP.lock().insert((tid, fd), handle);
132}
133
134/// Unregister a socket fd mapping
135pub fn unregister_socket_fd(tid: usize, fd: usize) {
136    FD_SOCKET_MAP.lock().remove(&(tid, fd));
137}
138
139/// Get socket handle from (tid, fd)
140pub fn get_socket_handle(tid: usize, fd: usize) -> Option<SocketHandle> {
141    FD_SOCKET_MAP.lock().get(&(tid, fd)).copied()
142}
143
144/// Update socket handle for an existing fd (used in accept)
145pub fn update_socket_handle(tid: usize, fd: usize, handle: SocketHandle) {
146    FD_SOCKET_MAP.lock().insert((tid, fd), handle);
147}
148
149/// Set local endpoint for a socket
150pub fn set_socket_local_endpoint(
151    file: &Arc<dyn crate::vfs::File>,
152    endpoint: IpEndpoint,
153) -> Result<(), ()> {
154    let any = file.as_any();
155    if let Some(socket_file) = any.downcast_ref::<SocketFile>() {
156        socket_file.set_local_endpoint(endpoint);
157        Ok(())
158    } else {
159        Err(())
160    }
161}
162
163/// Get local endpoint from a socket
164pub fn get_socket_local_endpoint(file: &Arc<dyn crate::vfs::File>) -> Option<IpEndpoint> {
165    let any = file.as_any();
166    any.downcast_ref::<SocketFile>()
167        .and_then(|socket_file| socket_file.get_local_endpoint())
168}
169
170/// Set remote endpoint for a socket
171pub fn set_socket_remote_endpoint(
172    file: &Arc<dyn crate::vfs::File>,
173    endpoint: IpEndpoint,
174) -> Result<(), ()> {
175    let any = file.as_any();
176    if let Some(socket_file) = any.downcast_ref::<SocketFile>() {
177        socket_file.set_remote_endpoint(endpoint);
178        Ok(())
179    } else {
180        Err(())
181    }
182}
183
184/// Get remote endpoint from a socket
185pub fn get_socket_remote_endpoint(file: &Arc<dyn crate::vfs::File>) -> Option<IpEndpoint> {
186    let any = file.as_any();
187    any.downcast_ref::<SocketFile>()
188        .and_then(|socket_file| socket_file.get_remote_endpoint())
189}
190
191/// Shutdown socket read
192pub fn socket_shutdown_read(file: &Arc<dyn crate::vfs::File>) {
193    if let Some(socket_file) = file.as_any().downcast_ref::<SocketFile>() {
194        socket_file.shutdown_read();
195    }
196}
197
198/// Shutdown socket write
199pub fn socket_shutdown_write(file: &Arc<dyn crate::vfs::File>) {
200    if let Some(socket_file) = file.as_any().downcast_ref::<SocketFile>() {
201        socket_file.shutdown_write();
202    }
203}
204
205/// Update socket handle in SocketFile (used in accept)
206pub fn update_socket_file_handle(
207    file: &Arc<dyn crate::vfs::File>,
208    new_handle: SocketHandle,
209) -> Result<(), ()> {
210    let any = file.as_any();
211    if let Some(socket_file) = any.downcast_ref::<SocketFile>() {
212        socket_file.set_handle(new_handle);
213        Ok(())
214    } else {
215        Err(())
216    }
217}
218
219/// Send data to a specific endpoint (for sendto syscall)
220pub fn socket_sendto(
221    handle: SocketHandle,
222    buf: &[u8],
223    endpoint: IpEndpoint,
224) -> Result<usize, FsError> {
225    let mut sockets = SOCKET_SET.lock();
226    match handle {
227        SocketHandle::Tcp(_) => Err(FsError::NotSupported), // TCP doesn't support sendto
228        SocketHandle::Udp(h) => {
229            let socket = sockets.get_mut::<udp::Socket>(h);
230            socket
231                .send_slice(buf, endpoint)
232                .map_err(|_| FsError::WouldBlock)?;
233            Ok(buf.len())
234        }
235    }
236}
237
238impl File for SocketFile {
239    fn readable(&self) -> bool {
240        let sockets = SOCKET_SET.lock();
241        match *self.handle.lock() {
242            SocketHandle::Tcp(h) => {
243                let socket = sockets.get::<tcp::Socket>(h);
244                socket.can_recv()
245            }
246            SocketHandle::Udp(h) => {
247                let socket = sockets.get::<udp::Socket>(h);
248                socket.can_recv()
249            }
250        }
251    }
252    fn writable(&self) -> bool {
253        let sockets = SOCKET_SET.lock();
254        match *self.handle.lock() {
255            SocketHandle::Tcp(h) => {
256                let socket = sockets.get::<tcp::Socket>(h);
257                socket.can_send()
258            }
259            SocketHandle::Udp(h) => {
260                let socket = sockets.get::<udp::Socket>(h);
261                socket.can_send()
262            }
263        }
264    }
265
266    fn read(&self, buf: &mut [u8]) -> Result<usize, FsError> {
267        if self.is_shutdown_read() {
268            return Ok(0); // EOF
269        }
270
271        let mut sockets = SOCKET_SET.lock();
272        let result = match *self.handle.lock() {
273            SocketHandle::Tcp(h) => {
274                let socket = sockets.get_mut::<tcp::Socket>(h);
275                socket.recv_slice(buf).map_err(|_| FsError::WouldBlock)
276            }
277            SocketHandle::Udp(h) => {
278                let socket = sockets.get_mut::<udp::Socket>(h);
279                socket
280                    .recv_slice(buf)
281                    .map(|(n, _)| n)
282                    .map_err(|_| FsError::WouldBlock)
283            }
284        };
285        if result.is_ok() {
286            crate::kernel::syscall::io::wake_poll_waiters();
287        }
288        result
289    }
290
291    fn write(&self, buf: &[u8]) -> Result<usize, FsError> {
292        if self.is_shutdown_write() {
293            return Err(FsError::BrokenPipe);
294        }
295
296        let mut sockets = SOCKET_SET.lock();
297        let result = match *self.handle.lock() {
298            SocketHandle::Tcp(h) => {
299                let socket = sockets.get_mut::<tcp::Socket>(h);
300                socket.send_slice(buf).map_err(|_| FsError::WouldBlock)
301            }
302            SocketHandle::Udp(h) => {
303                let endpoint = match self.get_remote_endpoint() {
304                    Some(ep) => ep,
305                    None => return Err(FsError::NotConnected),
306                };
307                let socket = sockets.get_mut::<udp::Socket>(h);
308                socket
309                    .send_slice(buf, endpoint)
310                    .map_err(|_| FsError::WouldBlock)?;
311                Ok(buf.len())
312            }
313        };
314        if result.is_ok() {
315            crate::kernel::syscall::io::wake_poll_waiters();
316        }
317        result
318    }
319
320    fn metadata(&self) -> Result<InodeMetadata, FsError> {
321        Err(FsError::NotSupported)
322    }
323
324    fn as_any(&self) -> &dyn core::any::Any {
325        self
326    }
327
328    fn flags(&self) -> OpenFlags {
329        *self.flags.lock()
330    }
331
332    fn set_status_flags(&self, new_flags: OpenFlags) -> Result<(), FsError> {
333        *self.flags.lock() = new_flags;
334        Ok(())
335    }
336
337    fn recvfrom(&self, buf: &mut [u8]) -> Result<(usize, Option<alloc::vec::Vec<u8>>), FsError> {
338        if self.is_shutdown_read() {
339            return Ok((0, None));
340        }
341
342        let mut sockets = SOCKET_SET.lock();
343        match *self.handle.lock() {
344            SocketHandle::Tcp(h) => {
345                let socket = sockets.get_mut::<tcp::Socket>(h);
346                let n = socket.recv_slice(buf).map_err(|_| FsError::WouldBlock)?;
347                let remote = socket.remote_endpoint().map(|ep| {
348                    let mut addr_buf = alloc::vec![0u8; 16];
349                    let _ = write_sockaddr_in_to_buf(&mut addr_buf, ep);
350                    addr_buf
351                });
352                Ok((n, remote))
353            }
354            SocketHandle::Udp(h) => {
355                let socket = sockets.get_mut::<udp::Socket>(h);
356                let (n, metadata) = socket.recv_slice(buf).map_err(|_| FsError::WouldBlock)?;
357                let mut addr_buf = alloc::vec![0u8; 16];
358                let _ = write_sockaddr_in_to_buf(&mut addr_buf, metadata.endpoint);
359                Ok((n, Some(addr_buf)))
360            }
361        }
362    }
363}
364
365pub fn create_tcp_socket() -> Result<SocketHandle, ()> {
366    // Allocate buffers with fallible allocation
367    let mut rx_vec = alloc::vec::Vec::new();
368    rx_vec.try_reserve(4096).map_err(|_| ())?;
369    rx_vec.resize(4096, 0);
370
371    let mut tx_vec = alloc::vec::Vec::new();
372    tx_vec.try_reserve(4096).map_err(|_| ())?;
373    tx_vec.resize(4096, 0);
374
375    let rx_buffer = tcp::SocketBuffer::new(rx_vec);
376    let tx_buffer = tcp::SocketBuffer::new(tx_vec);
377    let socket = tcp::Socket::new(rx_buffer, tx_buffer);
378    let handle = SOCKET_SET.lock().add(socket);
379    Ok(SocketHandle::Tcp(handle))
380}
381
382pub fn create_udp_socket() -> Result<SocketHandle, ()> {
383    // Allocate metadata buffers
384    let mut rx_meta_vec = alloc::vec::Vec::new();
385    rx_meta_vec.try_reserve(4).map_err(|_| ())?;
386    rx_meta_vec.resize(4, udp::PacketMetadata::EMPTY);
387
388    let mut tx_meta_vec = alloc::vec::Vec::new();
389    tx_meta_vec.try_reserve(4).map_err(|_| ())?;
390    tx_meta_vec.resize(4, udp::PacketMetadata::EMPTY);
391
392    // Allocate data buffers
393    let mut rx_data_vec = alloc::vec::Vec::new();
394    rx_data_vec.try_reserve(4096).map_err(|_| ())?;
395    rx_data_vec.resize(4096, 0);
396
397    let mut tx_data_vec = alloc::vec::Vec::new();
398    tx_data_vec.try_reserve(4096).map_err(|_| ())?;
399    tx_data_vec.resize(4096, 0);
400
401    let rx_buffer = udp::PacketBuffer::new(rx_meta_vec, rx_data_vec);
402    let tx_buffer = udp::PacketBuffer::new(tx_meta_vec, tx_data_vec);
403    let socket = udp::Socket::new(rx_buffer, tx_buffer);
404    let handle = SOCKET_SET.lock().add(socket);
405    Ok(SocketHandle::Udp(handle))
406}
407
408const AF_INET: u16 = 2;
409const SOCKADDR_IN_SIZE: usize = 16;
410
411/// Parse sockaddr_in structure
412pub fn parse_sockaddr_in(addr: *const u8, addrlen: u32) -> Result<IpEndpoint, ()> {
413    if (addrlen as usize) < SOCKADDR_IN_SIZE {
414        return Err(());
415    }
416
417    unsafe {
418        let family = core::ptr::read_unaligned(addr as *const u16);
419        if family != AF_INET {
420            return Err(());
421        }
422
423        let port = u16::from_be(core::ptr::read_unaligned(addr.add(2) as *const u16));
424        let ip_bytes = core::slice::from_raw_parts(addr.add(4), 4);
425        let ip = Ipv4Address::from_octets([ip_bytes[0], ip_bytes[1], ip_bytes[2], ip_bytes[3]]);
426
427        Ok(IpEndpoint::new(IpAddress::Ipv4(ip), port))
428    }
429}
430
431/// Write sockaddr_in to buffer
432fn write_sockaddr_in_to_buf(buf: &mut [u8], endpoint: IpEndpoint) -> Result<(), ()> {
433    if buf.len() < SOCKADDR_IN_SIZE {
434        return Err(());
435    }
436
437    // family
438    buf[0..2].copy_from_slice(&AF_INET.to_ne_bytes());
439
440    // port
441    buf[2..4].copy_from_slice(&endpoint.port.to_be_bytes());
442
443    // ip
444    match endpoint.addr {
445        IpAddress::Ipv4(ipv4) => {
446            buf[4..8].copy_from_slice(&ipv4.octets());
447        }
448        #[cfg(feature = "proto-ipv6")]
449        IpAddress::Ipv6(_) => {
450            return Err(()); // IPv6 not supported in AF_INET
451        }
452        _ => {
453            return Err(()); // Unknown address type
454        }
455    }
456
457    // zero padding
458    buf[8..16].fill(0);
459
460    Ok(())
461}
462
463/// Write sockaddr_in structure
464pub fn write_sockaddr_in(addr: *mut u8, addrlen: *mut u32, endpoint: IpEndpoint) -> Result<(), ()> {
465    if addr.is_null() || addrlen.is_null() {
466        return Ok(());
467    }
468
469    unsafe {
470        let len = *addrlen as usize;
471        if len < SOCKADDR_IN_SIZE {
472            return Err(());
473        }
474
475        let buf = core::slice::from_raw_parts_mut(addr, SOCKADDR_IN_SIZE);
476        write_sockaddr_in_to_buf(buf, endpoint)?;
477
478        *addrlen = SOCKADDR_IN_SIZE as u32;
479    }
480
481    Ok(())
482}
483
484/// Initialize global interface (should be called during network setup)
485pub fn init_global_interface(iface: Interface) {
486    *GLOBAL_INTERFACE.lock() = Some(iface);
487}
488
489/// Perform TCP connect with Context
490pub fn tcp_connect(handle: SmoltcpHandle, remote: IpEndpoint, local: IpEndpoint) -> Result<(), ()> {
491    let mut iface_guard = GLOBAL_INTERFACE.lock();
492    let iface = iface_guard.as_mut().ok_or(())?;
493
494    let mut sockets = SOCKET_SET.lock();
495    let socket = sockets.get_mut::<tcp::Socket>(handle);
496
497    socket
498        .connect(iface.context(), remote, local)
499        .map_err(|_| ())
500}
os/net/socket.rs

os/net/
socket.rs
