Socket::INet: recv

Signed-off-by: Andy-Python-Programmer <andypythonappdeveloper@gmail.com>

Author: Andy-Python-Programmer

src/aero_kernel/src/drivers/e1000.rs

@@ -592,7 +592,7 @@ impl Device {
     }
 
     fn handle_irq(&self) {
-        self.e1000.lock().handle_irq()
+        self.e1000.lock_irq().handle_irq()
     }
 }
 

src/aero_kernel/src/net/ethernet.rs

@@ -85,4 +85,13 @@ impl PacketHeader<Header> for Packet<Eth> {
             super::default_device().send(packet);
         }
     }
+
+    fn recv(&self) {
+        match self.header().typ {
+            Type::Ip => {
+                let packet: Packet<ip::Ipv4> = self.upgrade();
+                packet.recv()
+            }
+        }
+    }
 }

src/aero_kernel/src/net/ip.rs

@@ -63,6 +63,10 @@ impl Header {
     fn set_length(&mut self, length: u16) {
         self.length = BigEndian::from(length);
     }
+
+    fn length(&self) -> BigEndian<u16> {
+        self.length
+    }
 }
 
 #[derive(Clone)]
@@ -103,4 +107,17 @@ impl PacketHeader<Header> for Packet<Ipv4> {
     fn send(&self) {
         self.downgrade().send() // send the ethernet packet
     }
+
+    fn recv(&self) {
+        let mut packet = self.clone();
+        let header = self.header();
+
+        packet.len = header.length().to_native() as usize;
+        match header.protocol {
+            Type::Udp => {
+                let packet: Packet<udp::Udp> = packet.upgrade();
+                packet.recv();
+            }
+        }
+    }
 }

src/aero_kernel/src/net/mod.rs

@@ -79,6 +79,7 @@ pub trait PacketBaseTrait {
 pub trait PacketTrait: PacketBaseTrait {
     fn header_size(&self) -> usize;
 
+    // TODO: Rename as_slice{_mut} to payload{_mut}?
     fn as_slice_mut(&mut self) -> &mut [u8] {
         let hsize = self.header_size();
 
@@ -87,6 +88,15 @@ pub trait PacketTrait: PacketBaseTrait {
 
         unsafe { core::slice::from_raw_parts_mut(start.as_mut_ptr(), size) }
     }
+
+    fn as_slice(&self) -> &[u8] {
+        let hsize = self.header_size();
+
+        let start = self.addr() + hsize;
+        let size = self.len() - hsize;
+
+        unsafe { core::slice::from_raw_parts(start.as_ptr(), size) }
+    }
 }
 
 impl<T: ConstPacketKind> PacketTrait for Packet<T> {
@@ -138,6 +148,7 @@ pub trait PacketDownHierarchy<B: ConstPacketKind>: PacketBaseTrait {
 
 pub trait PacketHeader<H>: PacketBaseTrait {
     fn send(&self);
+    fn recv(&self);
 
     fn header(&self) -> &H {
         self.addr().read_mut::<H>().unwrap()
@@ -155,9 +166,8 @@ fn packet_processor_thread() {
     let device = default_device();
 
     loop {
-        log::debug!("bruh!");
         let packet = device.recv();
-        log::debug!("{packet:?}");
+        packet.packet.recv();
     }
 }
 

src/aero_kernel/src/net/udp.rs

@@ -62,6 +62,18 @@ impl PacketHeader<Header> for Packet<Udp> {
 
         self.downgrade().send() // send the IP packet
     }
+
+    fn recv(&self) {
+        let header = self.header();
+        let dest_port = header.dst_port().to_native();
+
+        let handlers = HANDLERS.read();
+        let handler = handlers
+            .get(&dest_port)
+            .expect("udp: no handler registered");
+
+        handler.recv(self.clone())
+    }
 }
 
 #[repr(C, packed)]
@@ -82,11 +94,17 @@ impl Header {
             checksum::calculate_with_len(self, length.to_native() as usize),
         ]);
     }
+
+    fn dst_port(&self) -> BigEndian<u16> {
+        self.dst_port
+    }
 }
 
 static HANDLERS: RwLock<BTreeMap<u16, Arc<dyn UdpHandler>>> = RwLock::new(BTreeMap::new());
 
-pub trait UdpHandler: Send + Sync {}
+pub trait UdpHandler: Send + Sync {
+    fn recv(&self, packet: Packet<Udp>);
+}
 
 pub fn alloc_ephemeral_port(socket: Arc<dyn UdpHandler>) -> Option<u16> {
     const EPHEMERAL_START: u16 = 49152;

src/aero_kernel/src/socket/inet.rs

@@ -17,16 +17,20 @@
  * along with Aero. If not, see <https://www.gnu.org/licenses/>.
  */
 
-use aero_syscall::{IpProtocol, SocketAddrInet, SocketType, SyscallError};
+use aero_syscall::socket::MessageHeader;
+use aero_syscall::{IpProtocol, OpenFlags, SocketAddrInet, SocketType, SyscallError};
 use alloc::sync::{Arc, Weak};
 use alloc::vec::Vec;
+use spin::Once;
 
-use crate::fs::inode::{FileType, INodeInterface, Metadata};
+use crate::fs::cache::DirCacheItem;
+use crate::fs::file_table::FileHandle;
+use crate::fs::inode::{FileType, INodeInterface, Metadata, PollFlags};
 use crate::fs::{self, FileSystemError};
 use crate::net::ip::Ipv4Addr;
 use crate::net::udp::{self, Udp, UdpHandler};
 use crate::net::{Packet, PacketHeader, PacketTrait};
-use crate::utils::sync::Mutex;
+use crate::utils::sync::{BlockQueue, Mutex};
 
 #[derive(Default)]
 enum SocketState {
@@ -41,22 +45,32 @@ struct InetSocketInner {
     /// The address that the socket has been bound to.
     address: Option<SocketAddrInet>,
     state: SocketState,
+    incoming: Vec<Packet<Udp>>,
 }
 
 pub struct InetSocket {
     typ: SocketType,
     protocol: IpProtocol,
     inner: Mutex<InetSocketInner>,
+    wq: BlockQueue,
+    handle: Once<Arc<FileHandle>>,
 
     sref: Weak<Self>,
 }
 
 impl InetSocket {
     pub fn new(typ: SocketType, protocol: IpProtocol) -> Result<Arc<Self>, SyscallError> {
+        if typ != SocketType::Dgram && protocol != IpProtocol::Udp {
+            return Err(SyscallError::EINVAL);
+        }
+
         Ok(Arc::new_cyclic(|sref| Self {
             typ,
             protocol,
 
+            wq: BlockQueue::new(),
+            handle: Once::new(),
+
             inner: Mutex::new(Default::default()),
             sref: sref.clone(),
         }))
@@ -88,9 +102,27 @@ impl InetSocket {
             _ => unreachable!(),
         }
     }
+
+    pub fn is_non_block(&self) -> bool {
+        self.handle
+            .get()
+            .expect("inet: not bound to an fd")
+            .flags
+            .read()
+            .contains(OpenFlags::O_NONBLOCK)
+    }
 }
 
 impl INodeInterface for InetSocket {
+    fn open(
+        &self,
+        _flags: aero_syscall::OpenFlags,
+        handle: Arc<FileHandle>,
+    ) -> fs::Result<Option<DirCacheItem>> {
+        self.handle.call_once(|| handle);
+        Ok(None)
+    }
+
     fn metadata(&self) -> fs::Result<fs::inode::Metadata> {
         Ok(Metadata {
             id: 0,
@@ -130,7 +162,7 @@ impl INodeInterface for InetSocket {
         }
     }
 
-    fn send(&self, message_hdr: &mut aero_syscall::socket::MessageHeader) -> fs::Result<usize> {
+    fn send(&self, message_hdr: &mut MessageHeader) -> fs::Result<usize> {
         let name = message_hdr
             .name_mut::<SocketAddrInet>()
             .cloned()
@@ -139,9 +171,17 @@ impl INodeInterface for InetSocket {
         let dest_port = name.port.to_native();
         let dest_ip = Ipv4Addr::new(name.addr());
 
-        let src_port = self.src_port().unwrap_or_else(|| {
-            udp::alloc_ephemeral_port(self.sref()).expect("inet: out of ephemeral ports")
-        });
+        let src_port;
+
+        if let Some(port) = self.src_port() {
+            src_port = port;
+        } else {
+            src_port = udp::alloc_ephemeral_port(self.sref()).ok_or(FileSystemError::WouldBlock)?;
+            log::debug!("Inet::send(): allocated ephemeral port {}", src_port);
+
+            // FIXME: handle ephemeral port INET socket
+            return Err(FileSystemError::NotSupported);
+        }
 
         let data = message_hdr
             .iovecs()
@@ -159,6 +199,41 @@ impl INodeInterface for InetSocket {
         packet.send();
         Ok(data.len())
     }
+
+    fn recv(&self, message_hdr: &mut MessageHeader) -> fs::Result<usize> {
+        if self.inner.lock_irq().incoming.is_empty() && self.is_non_block() {
+            return Err(FileSystemError::WouldBlock);
+        }
+
+        let mut this = self.wq.block_on(&self.inner, |e| !e.incoming.is_empty())?;
+        let packet = this.incoming.pop().expect("recv: someone was greedy");
+
+        let mut data = packet.as_slice().to_vec();
+
+        Ok(message_hdr
+            .iovecs_mut()
+            .iter_mut()
+            .map(|iovec| {
+                let iovec = iovec.as_slice_mut();
+                let size = core::cmp::min(iovec.len(), data.len());
+                iovec.copy_from_slice(&data.drain(..size).collect::<Vec<_>>());
+                size
+            })
+            .sum::<usize>())
+    }
+
+    fn poll(&self, table: Option<&mut fs::inode::PollTable>) -> fs::Result<PollFlags> {
+        if let Some(table) = table {
+            table.insert(&self.wq);
+        }
+
+        Ok(PollFlags::OUT)
+    }
 }
 
-impl UdpHandler for InetSocket {}
+impl UdpHandler for InetSocket {
+    fn recv(&self, packet: Packet<Udp>) {
+        self.inner.lock_irq().incoming.push(packet);
+        self.wq.notify_complete();
+    }
+}

src/aero_kernel/src/syscall/fs.rs

@@ -381,7 +381,7 @@ pub fn fcntl(fd: usize, command: usize, arg: usize) -> Result<usize, SyscallErro
 
         aero_syscall::prelude::F_SETFL => {
             let flags = OpenFlags::from_bits_truncate(arg);
-            handle.flags.write().insert(flags);
+            *handle.flags.write() = flags;
 
             Ok(0)
         }

src/aero_kernel/src/syscall/net.rs

@@ -124,7 +124,7 @@ fn create_socket(
 
     let socket = match domain as u32 {
         AF_UNIX => UnixSocket::new() as Arc<dyn INodeInterface>,
-        PF_INET => InetSocket::new(typ, protocol)? as Arc<dyn INodeInterface>,
+        AF_INET => InetSocket::new(typ, protocol)? as Arc<dyn INodeInterface>,
 
         _ => {
             log::warn!(