diff --git a/litebox/src/net/errors.rs b/litebox/src/net/errors.rs index e15e6db8b..028112c39 100644 --- a/litebox/src/net/errors.rs +++ b/litebox/src/net/errors.rs @@ -109,3 +109,22 @@ pub enum SetTcpOptionError { #[error("Not a TCP socket")] NotTcpSocket, } + +/// Possible errors from [`Network::with_metadata`] and [`Network::with_metadata_mut`] +#[non_exhaustive] +#[derive(Error, Debug)] +pub enum MetadataError { + #[error("Not a valid open file descriptor")] + InvalidFd, + #[error("no such metadata available")] + NoSuchMetadata, +} + +/// Possible errors from [`Network::set_socket_metadata`] +#[non_exhaustive] +#[derive(Error, Debug)] +pub enum SetMetadataError { + #[error("Not a valid open file descriptor")] + // Note: we return the T just so we are not dropping data + InvalidFd(T), +} diff --git a/litebox/src/net/mod.rs b/litebox/src/net/mod.rs index 756a479dd..15d815292 100644 --- a/litebox/src/net/mod.rs +++ b/litebox/src/net/mod.rs @@ -7,6 +7,7 @@ use core::net::{Ipv4Addr, SocketAddr}; use crate::event::Events; use crate::fd::InternalFd; use crate::platform::Instant; +use crate::utilities::anymap::AnyMap; use crate::{LiteBox, platform, sync}; use crate::{event::EventManager, fd::SocketFd}; @@ -21,8 +22,8 @@ mod phy; mod tests; use errors::{ - AcceptError, BindError, CloseError, ConnectError, ListenError, ReceiveError, SendError, - SocketError, + AcceptError, BindError, CloseError, ConnectError, ListenError, MetadataError, ReceiveError, + SendError, SetMetadataError, SocketError, }; use local_ports::{LocalPort, LocalPortAllocator}; @@ -134,6 +135,11 @@ struct SocketHandle { handle: smoltcp::iface::SocketHandle, // Protocol-specific data specific: ProtocolSpecific, + /// Socket-level metadata + /// TODO: FD-specific metadata (similar to fs) will be added in the future. + /// This needs to be handled when switching to the Arc-based implementation + /// (possibly connected to #31) and is tracked in #120. + socket_metadata: AnyMap, } impl core::ops::Deref for SocketHandle { @@ -541,6 +547,7 @@ where Protocol::Icmp => unimplemented!(), Protocol::Raw { protocol } => unimplemented!(), }, + socket_metadata: AnyMap::new(), })) } @@ -807,6 +814,7 @@ where local_port, server_socket: None, }), + socket_metadata: AnyMap::new(), })) } ProtocolSpecific::Udp(_) => unimplemented!(), @@ -907,6 +915,55 @@ where } } } + + /// Apply `f` on metadata at a socket fd, if it exists. + pub fn with_metadata( + &self, + fd: &SocketFd, + f: impl FnOnce(&T) -> R, + ) -> Result { + let socket_handle = self.handles[fd.x.as_usize()] + .as_ref() + .ok_or(MetadataError::InvalidFd)?; + + if let Some(m) = socket_handle.socket_metadata.get::() { + Ok(f(m)) + } else { + Err(MetadataError::NoSuchMetadata) + } + } + + /// Similar to [`Self::with_metadata`] but mutable. + pub fn with_metadata_mut( + &mut self, + fd: &SocketFd, + f: impl FnOnce(&mut T) -> R, + ) -> Result { + let socket_handle = self.handles[fd.x.as_usize()] + .as_mut() + .ok_or(MetadataError::InvalidFd)?; + + if let Some(m) = socket_handle.socket_metadata.get_mut::() { + Ok(f(m)) + } else { + Err(MetadataError::NoSuchMetadata) + } + } + + /// Store arbitrary metadata into a socket. + /// + /// Returns the old metadata if any such metadata exists. + pub fn set_socket_metadata( + &mut self, + fd: &SocketFd, + metadata: T, + ) -> Result, SetMetadataError> { + let Some(socket_handle) = self.handles[fd.x.as_usize()].as_mut() else { + return Err(SetMetadataError::InvalidFd(metadata)); + }; + + Ok(socket_handle.socket_metadata.insert(metadata)) + } } /// Protocols for sockets supported by LiteBox diff --git a/litebox/src/net/tests.rs b/litebox/src/net/tests.rs index cf4746d14..d9b4158c4 100644 --- a/litebox/src/net/tests.rs +++ b/litebox/src/net/tests.rs @@ -2,6 +2,7 @@ use platform::mock::MockPlatform; use super::*; +use alloc::string::String; use core::net::SocketAddrV4; use core::str::FromStr; @@ -97,3 +98,68 @@ fn test_bidirectional_tcp_communication_automatic() { network.set_platform_interaction(PlatformInteraction::Automatic); bidi_tcp_comms(network, |_| {}); } + +#[test] +fn test_socket_metadata() { + let litebox = LiteBox::new(MockPlatform::new()); + let mut network = Network::new(&litebox); + + // Create a socket + let socket_fd = network + .socket(Protocol::Tcp) + .expect("Failed to create TCP socket"); + + // Test setting and getting socket metadata + let test_data = String::from("socket-level metadata"); + let old_metadata = network + .set_socket_metadata(&socket_fd, test_data.clone()) + .expect("Failed to set socket metadata"); + assert!(old_metadata.is_none(), "Expected no previous metadata"); + + // Test reading socket metadata + let retrieved_data = network + .with_metadata(&socket_fd, |data: &String| data.clone()) + .expect("Failed to get socket metadata"); + assert_eq!(retrieved_data, test_data, "Retrieved metadata should match"); + + // Test overwriting socket metadata + let new_test_data = String::from("updated socket metadata"); + let old_metadata = network + .set_socket_metadata(&socket_fd, new_test_data.clone()) + .expect("Failed to update socket metadata"); + assert_eq!( + old_metadata, + Some(test_data), + "Should return previous metadata" + ); + + // Test retrieving updated metadata + let updated_data = network + .with_metadata(&socket_fd, |data: &String| data.clone()) + .expect("Failed to get updated socket metadata"); + assert_eq!(updated_data, new_test_data, "Updated metadata should match"); + + // Close the socket + network.close(socket_fd).expect("Failed to close socket"); +} + +#[test] +fn test_metadata_errors() { + let litebox = LiteBox::new(MockPlatform::new()); + let mut network = Network::new(&litebox); + + // Create a socket + let socket_fd = network + .socket(Protocol::Tcp) + .expect("Failed to create TCP socket"); + + // Test nonexistent metadata type + let result = network.with_metadata(&socket_fd, |_data: &String| ()); + assert!( + matches!(result, Err(MetadataError::NoSuchMetadata)), + "Should return NoSuchMetadata error" + ); + + // Close the socket + network.close(socket_fd).expect("Failed to close socket"); +}