kernel/network/
udp.rs

1//! UDP protocol layer
2//!
3//! This module provides UDP datagram handling for the network stack.
4//! It implements the NetworkLayer trait and provides UDP socket functionality.
5
6use alloc::collections::BTreeMap;
7use alloc::string::String;
8use alloc::sync::{Arc, Weak};
9use alloc::vec::Vec;
10use spin::{Mutex, RwLock};
11
12use crate::early_println;
13use crate::network::Ipv4Address;
14use crate::network::protocol_stack::get_network_manager;
15use crate::network::protocol_stack::{LayerContext, NetworkLayer, NetworkLayerStats, SocketConfig};
16use crate::network::socket::{
17    Inet4SocketAddress, SocketAddress, SocketControl, SocketError, SocketObject, SocketProtocol,
18    SocketState, SocketType,
19};
20use crate::object::capability::selectable::Selectable;
21
22/// Helper function to get local IP address bytes from the default interface
23fn get_local_ip_bytes() -> [u8; 4] {
24    let manager = get_network_manager();
25    if let Some(default_iface) = manager.get_default_interface() {
26        if let Some(ip_layer) = manager.get_layer("ip") {
27            if let Some(ipv4_layer) = ip_layer
28                .as_any()
29                .downcast_ref::<crate::network::ipv4::Ipv4Layer>()
30            {
31                if let Some(addr) = ipv4_layer.get_primary_ip(default_iface.name()) {
32                    return addr.as_bytes();
33                }
34            }
35        }
36    }
37    [0u8; 4]
38}
39
40/// UDP header (8 bytes)
41#[derive(Debug, Clone, Copy)]
42#[repr(C, packed)]
43pub struct UdpHeader {
44    /// Source port
45    pub src_port: u16,
46    /// Destination port
47    pub dst_port: u16,
48    /// Length (header + data)
49    pub length: u16,
50    /// Checksum
51    pub checksum: u16,
52}
53
54impl UdpHeader {
55    /// Create a new UDP header
56    pub fn new(src_port: u16, dst_port: u16, length: u16) -> Self {
57        Self {
58            src_port,
59            dst_port,
60            length,
61            checksum: 0,
62        }
63    }
64
65    /// Calculate UDP checksum
66    pub fn calculate_checksum(&self, src_ip: [u8; 4], dst_ip: [u8; 4], data: &[u8]) -> u16 {
67        let mut sum: u32 = 0;
68
69        // Pseudo-header: src IP (4) + dst IP (4) + zero (1) + protocol (1) + UDP length (2)
70        sum += u16::from_be_bytes([src_ip[0], src_ip[1]]) as u32;
71        sum = (sum & 0xFFFF) + (sum >> 16);
72        sum += u16::from_be_bytes([src_ip[2], src_ip[3]]) as u32;
73        sum = (sum & 0xFFFF) + (sum >> 16);
74        sum += u16::from_be_bytes([dst_ip[0], dst_ip[1]]) as u32;
75        sum = (sum & 0xFFFF) + (sum >> 16);
76        sum += u16::from_be_bytes([dst_ip[2], dst_ip[3]]) as u32;
77        sum = (sum & 0xFFFF) + (sum >> 16);
78        sum += 17u32; // UDP protocol number
79        sum = (sum & 0xFFFF) + (sum >> 16);
80        sum += self.length as u32;
81        sum = (sum & 0xFFFF) + (sum >> 16);
82
83        // UDP header (except checksum field)
84        sum += self.src_port as u32;
85        sum = (sum & 0xFFFF) + (sum >> 16);
86        sum += self.dst_port as u32;
87        sum = (sum & 0xFFFF) + (sum >> 16);
88        sum += self.length as u32;
89        sum = (sum & 0xFFFF) + (sum >> 16);
90
91        // Data
92        for chunk in data.chunks(2) {
93            if chunk.len() == 2 {
94                sum += u16::from_be_bytes([chunk[0], chunk[1]]) as u32;
95                sum = (sum & 0xFFFF) + (sum >> 16);
96            } else if chunk.len() == 1 {
97                sum += (chunk[0] as u32) << 8;
98                sum = (sum & 0xFFFF) + (sum >> 16);
99            }
100        }
101
102        // Final carry
103        while sum >> 16 != 0 {
104            sum = (sum & 0xFFFF) + (sum >> 16);
105        }
106
107        !sum as u16
108    }
109
110    /// Serialize header to bytes
111    pub fn to_bytes(&self) -> [u8; 8] {
112        let mut bytes = [0u8; 8];
113        bytes[0..2].copy_from_slice(&self.src_port.to_be_bytes());
114        bytes[2..4].copy_from_slice(&self.dst_port.to_be_bytes());
115        bytes[4..6].copy_from_slice(&self.length.to_be_bytes());
116        bytes[6..8].copy_from_slice(&self.checksum.to_be_bytes());
117        bytes
118    }
119
120    /// Parse header from bytes
121    pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
122        if bytes.len() < 8 {
123            return None;
124        }
125
126        Some(Self {
127            src_port: u16::from_be_bytes([bytes[0], bytes[1]]),
128            dst_port: u16::from_be_bytes([bytes[2], bytes[3]]),
129            length: u16::from_be_bytes([bytes[4], bytes[5]]),
130            checksum: u16::from_be_bytes([bytes[6], bytes[7]]),
131        })
132    }
133}
134
135/// UDP socket
136///
137/// Implements SocketObject for UDP datagram communication.
138pub struct UdpSocket {
139    /// Local address
140    local_addr: RwLock<Option<SocketAddress>>,
141    /// Remote address (for connected sockets)
142    remote_addr: RwLock<Option<SocketAddress>>,
143    /// Send buffer
144    send_buffer: Mutex<Vec<Vec<u8>>>,
145    /// Receive buffer
146    recv_buffer: Mutex<Vec<Vec<u8>>>,
147    /// Socket state
148    state: RwLock<SocketState>,
149    /// Reference to UDP layer
150    udp_layer: Arc<UdpLayer>,
151    /// Weak self reference for registration
152    self_weak: Weak<UdpSocket>,
153    /// Receive waker for blocking I/O
154    recv_waker: Mutex<Option<alloc::sync::Arc<crate::sync::Waker>>>,
155    /// Send waker for blocking I/O
156    send_waker: Mutex<Option<alloc::sync::Arc<crate::sync::Waker>>>,
157    /// Blocking mode (default: true)
158    blocking_mode: spin::Mutex<bool>,
159}
160
161impl UdpSocket {
162    /// Create a new UDP socket
163    pub fn new(udp_layer: Arc<UdpLayer>) -> Arc<Self> {
164        Arc::new_cyclic(|weak| Self {
165            local_addr: RwLock::new(None),
166            remote_addr: RwLock::new(None),
167            send_buffer: Mutex::new(Vec::new()),
168            recv_buffer: Mutex::new(Vec::new()),
169            state: RwLock::new(SocketState::Unconnected),
170            udp_layer,
171            self_weak: weak.clone(),
172            recv_waker: Mutex::new(None),
173            send_waker: Mutex::new(None),
174            blocking_mode: spin::Mutex::new(true),
175        })
176    }
177
178    /// Deliver received datagram to this socket
179    pub fn deliver_datagram(&self, data: Vec<u8>) {
180        self.recv_buffer.lock().push(data);
181        // Wake up any waiting reader
182        if let Some(waker) = self.recv_waker.lock().as_ref() {
183            waker.wake_one();
184        }
185    }
186}
187
188impl SocketObject for UdpSocket {
189    fn socket_type(&self) -> SocketType {
190        SocketType::Datagram
191    }
192
193    fn socket_domain(&self) -> crate::network::socket::SocketDomain {
194        crate::network::socket::SocketDomain::Inet4
195    }
196
197    fn socket_protocol(&self) -> SocketProtocol {
198        SocketProtocol::Udp
199    }
200
201    fn as_any(&self) -> &dyn core::any::Any {
202        self
203    }
204
205    fn sendto(
206        &self,
207        data: &[u8],
208        address: &SocketAddress,
209        _flags: u32,
210    ) -> Result<usize, SocketError> {
211        match address {
212            SocketAddress::Inet(inet) => {
213                let addr = inet.addr;
214                let port = inet.port;
215                // Queue the datagram for sending
216                let mut buffer = self.send_buffer.lock();
217                let datagram = data.to_vec();
218                buffer.push(datagram.clone());
219
220                // Update state
221                *self.remote_addr.write() = Some(address.clone());
222                *self.state.write() = SocketState::Connected;
223
224                // Try to send through UDP layer
225                self.udp_layer.send_datagram(self, addr, port, datagram)?;
226
227                Ok(data.len())
228            }
229            _ => Err(SocketError::InvalidAddress),
230        }
231    }
232
233    fn recvfrom(
234        &self,
235        buffer: &mut [u8],
236        _flags: u32,
237    ) -> Result<(usize, SocketAddress), SocketError> {
238        // Blocking mode: wait for data
239        if !self.is_nonblocking() {
240            let interest = crate::object::capability::selectable::ReadyInterest {
241                read: true,
242                write: false,
243                except: false,
244            };
245            let task = crate::task::mytask();
246            if let Some(t) = task {
247                let trapframe = t.get_trapframe();
248                Selectable::wait_until_ready(self, interest, trapframe, None);
249            }
250        }
251
252        let mut recv_buf = self.recv_buffer.lock();
253
254        if recv_buf.is_empty() {
255            return Err(SocketError::WouldBlock);
256        }
257
258        let datagram = recv_buf.remove(0);
259        let len = buffer.len().min(datagram.len());
260        buffer[..len].copy_from_slice(&datagram[..len]);
261
262        Ok((
263            len,
264            self.remote_addr
265                .read()
266                .clone()
267                .unwrap_or(SocketAddress::Unspecified),
268        ))
269    }
270}
271
272impl SocketControl for UdpSocket {
273    fn bind(&self, address: &SocketAddress) -> Result<(), SocketError> {
274        match address {
275            SocketAddress::Inet(inet) => {
276                let addr = inet.addr;
277                let port = inet.port;
278                let mut config = SocketConfig::new();
279                config.set("udp_local_port", &port.to_be_bytes());
280                config.set("ip_local", &addr);
281
282                // Configure UDP layer
283                self.udp_layer
284                    .configure_socket(self.self_weak.clone(), &config)?;
285
286                *self.local_addr.write() = Some(address.clone());
287                *self.state.write() = SocketState::Bound;
288                Ok(())
289            }
290            _ => Err(SocketError::InvalidAddress),
291        }
292    }
293
294    fn connect(&self, address: &SocketAddress) -> Result<(), SocketError> {
295        match address {
296            SocketAddress::Inet(_) => {
297                *self.remote_addr.write() = Some(address.clone());
298                *self.state.write() = SocketState::Connected;
299                Ok(())
300            }
301            _ => Err(SocketError::InvalidAddress),
302        }
303    }
304
305    fn listen(&self, _backlog: usize) -> Result<(), SocketError> {
306        Err(SocketError::NotSupported)
307    }
308
309    fn accept(&self) -> Result<Arc<dyn SocketObject>, SocketError> {
310        Err(SocketError::NotSupported)
311    }
312
313    fn getpeername(&self) -> Result<SocketAddress, SocketError> {
314        self.remote_addr
315            .read()
316            .clone()
317            .ok_or(SocketError::NotConnected)
318    }
319
320    fn getsockname(&self) -> Result<SocketAddress, SocketError> {
321        self.local_addr
322            .read()
323            .clone()
324            .ok_or(SocketError::InvalidAddress)
325    }
326
327    fn shutdown(&self, _how: crate::network::socket::ShutdownHow) -> Result<(), SocketError> {
328        *self.state.write() = SocketState::Closed;
329        Ok(())
330    }
331
332    fn is_connected(&self) -> bool {
333        *self.state.read() == SocketState::Connected
334    }
335
336    fn state(&self) -> SocketState {
337        *self.state.read()
338    }
339}
340
341impl crate::ipc::StreamIpcOps for UdpSocket {
342    fn is_connected(&self) -> bool {
343        SocketControl::is_connected(self)
344    }
345
346    fn peer_count(&self) -> usize {
347        if SocketControl::is_connected(self) {
348            1
349        } else {
350            0
351        }
352    }
353
354    fn description(&self) -> String {
355        alloc::format!("UDP socket")
356    }
357}
358
359impl crate::object::capability::StreamOps for UdpSocket {
360    fn read(&self, buffer: &mut [u8]) -> Result<usize, crate::object::capability::StreamError> {
361        use crate::object::capability::selectable::Selectable;
362
363        if !Selectable::is_nonblocking(self) {
364            // Blocking mode: wait for data
365            let task = crate::task::mytask();
366            if let Some(t) = task {
367                let trapframe = t.get_trapframe();
368                let interest = crate::object::capability::selectable::ReadyInterest {
369                    read: true,
370                    write: false,
371                    except: false,
372                };
373                Selectable::wait_until_ready(self, interest, trapframe, None);
374            }
375        }
376
377        let (len, _) = self.recvfrom(buffer, 0).map_err(|e| match e {
378            SocketError::WouldBlock => crate::object::capability::StreamError::WouldBlock,
379            _ => crate::object::capability::StreamError::Other("udp recv error".into()),
380        })?;
381        Ok(len)
382    }
383
384    fn write(&self, data: &[u8]) -> Result<usize, crate::object::capability::StreamError> {
385        let remote_addr = self
386            .remote_addr
387            .read()
388            .clone()
389            .unwrap_or(SocketAddress::Unspecified);
390        self.sendto(data, &remote_addr, 0)
391            .map_err(|_| crate::object::capability::StreamError::Other("udp send error".into()))?;
392        Ok(data.len())
393    }
394}
395
396impl crate::object::capability::CloneOps for UdpSocket {
397    fn custom_clone(&self) -> crate::object::KernelObject {
398        crate::object::KernelObject::Socket(UdpSocket::new(self.udp_layer.clone()))
399    }
400}
401
402impl crate::object::capability::Selectable for UdpSocket {
403    fn current_ready(
404        &self,
405        interest: crate::object::capability::selectable::ReadyInterest,
406    ) -> crate::object::capability::selectable::ReadySet {
407        let mut ready = crate::object::capability::selectable::ReadySet::none();
408
409        if interest.read {
410            let recv_buf = self.recv_buffer.lock();
411            ready.read = !recv_buf.is_empty();
412        }
413
414        if interest.write {
415            // UDP writes are always ready (no congestion control)
416            ready.write = true;
417        }
418
419        ready
420    }
421
422    fn wait_until_ready(
423        &self,
424        interest: crate::object::capability::selectable::ReadyInterest,
425        trapframe: &mut crate::arch::Trapframe,
426        timeout_ticks: Option<u64>,
427    ) -> crate::object::capability::selectable::SelectWaitOutcome {
428        let current = self.current_ready(interest);
429        if (interest.read && current.read) || (interest.write && current.write) {
430            return crate::object::capability::selectable::SelectWaitOutcome::Ready;
431        }
432
433        let task_id = {
434            use crate::arch::get_cpu;
435            use crate::sched::scheduler::get_scheduler;
436            let cpu_id = get_cpu().get_cpuid();
437            get_scheduler().get_current_task_id(cpu_id).unwrap_or(0)
438        };
439
440        let woke = if interest.read {
441            let waker = {
442                let mut waker_lock = self.recv_waker.lock();
443                waker_lock
444                    .get_or_insert_with(|| {
445                        alloc::sync::Arc::new(crate::sync::Waker::new_interruptible("udp_recv"))
446                    })
447                    .clone()
448            };
449            waker.wait_with_timeout(task_id, trapframe, timeout_ticks)
450        } else {
451            true
452        };
453
454        if timeout_ticks.is_some() && !woke {
455            let after = self.current_ready(interest);
456            if (interest.read && !after.read) && (interest.write && !after.write) {
457                return crate::object::capability::selectable::SelectWaitOutcome::TimedOut;
458            }
459        }
460
461        crate::object::capability::selectable::SelectWaitOutcome::Ready
462    }
463
464    fn set_nonblocking(&self, enabled: bool) {
465        *self.blocking_mode.lock() = !enabled;
466    }
467
468    fn is_nonblocking(&self) -> bool {
469        !*self.blocking_mode.lock()
470    }
471}
472
473/// UDP layer
474///
475/// Manages UDP port bindings and handles UDP datagrams.
476pub struct UdpLayer {
477    /// Port-to-socket mapping for receiving datagrams
478    port_map: RwLock<BTreeMap<u16, alloc::sync::Weak<UdpSocket>>>,
479    /// Port allocation (ephemeral ports start from 49152)
480    next_ephemeral_port: Mutex<u16>,
481    /// Statistics
482    stats: RwLock<NetworkLayerStats>,
483    self_weak: Weak<UdpLayer>,
484}
485
486impl UdpLayer {
487    /// Create a new UDP layer
488    pub fn new() -> Arc<Self> {
489        Arc::new_cyclic(|weak| Self {
490            port_map: RwLock::new(BTreeMap::new()),
491            next_ephemeral_port: Mutex::new(49152),
492            stats: RwLock::new(NetworkLayerStats::default()),
493            self_weak: weak.clone(),
494        })
495    }
496
497    /// Initialize and register the UDP layer with NetworkManager
498    ///
499    /// Registers with NetworkManager and registers itself with Ipv4Layer
500    /// for protocol number 17 (UDP).
501    ///
502    /// # Panics
503    ///
504    /// Panics if Ipv4Layer is not registered (must be initialized first).
505    pub fn init(network_manager: &crate::network::NetworkManager) {
506        let layer = Self::new();
507        network_manager.register_layer("udp", layer.clone());
508
509        // Register with IPv4 layer for UDP packets (protocol 17)
510        let ipv4 = network_manager
511            .get_layer("ip")
512            .expect("Ipv4Layer must be initialized before UdpLayer");
513        ipv4.register_protocol(crate::network::ipv4::protocol::UDP as u16, layer);
514    }
515
516    /// Create a new UDP socket
517    pub fn create_socket(&self) -> Arc<UdpSocket> {
518        let layer = self
519            .self_weak
520            .upgrade()
521            .expect("udp layer is not initialized");
522        UdpSocket::new(layer)
523    }
524
525    /// Allocate an ephemeral port
526    pub fn allocate_port(&self) -> u16 {
527        let mut next_port = self.next_ephemeral_port.lock();
528        let port = *next_port;
529
530        *next_port = if port == 65535 { 49152 } else { port + 1 };
531
532        port
533    }
534
535    /// Register a socket for a specific port
536    pub fn register_port(&self, port: u16, socket: alloc::sync::Weak<UdpSocket>) {
537        self.port_map.write().insert(port, socket);
538    }
539
540    /// Unregister a specific socket from a port
541    ///
542    /// Only removes the port entry if the registered socket matches.
543    pub fn unregister_socket(&self, port: u16, socket: &alloc::sync::Weak<UdpSocket>) {
544        let mut map = self.port_map.write();
545        if let Some(existing) = map.get(&port) {
546            if existing.ptr_eq(socket) {
547                map.remove(&port);
548            }
549        }
550    }
551
552    /// Find socket for a destination port
553    pub fn find_socket(&self, port: u16) -> Option<Arc<UdpSocket>> {
554        self.port_map
555            .read()
556            .get(&port)
557            .and_then(|weak| weak.upgrade())
558    }
559
560    /// Configure a UDP socket (bind)
561    pub fn configure_socket(
562        &self,
563        socket: Weak<UdpSocket>,
564        config: &SocketConfig,
565    ) -> Result<(), SocketError> {
566        let port = config
567            .get_u16("udp_local_port")
568            .ok_or(SocketError::InvalidAddress)?;
569
570        // Register the port
571        self.register_port(port, socket);
572
573        // TODO: Configure IP layer with local address
574        Ok(())
575    }
576
577    /// Send a UDP datagram
578    pub fn send_datagram(
579        &self,
580        socket: &UdpSocket,
581        dest_ip: [u8; 4],
582        dest_port: u16,
583        data: Vec<u8>,
584    ) -> Result<(), SocketError> {
585        let (src_ip_bytes, src_port) = match socket.local_addr.read().clone() {
586            Some(SocketAddress::Inet(inet)) => {
587                if inet.addr == [0, 0, 0, 0] {
588                    let ip = get_local_ip_bytes();
589                    (ip, inet.port)
590                } else {
591                    (inet.addr, inet.port)
592                }
593            }
594            _ => {
595                // Get local IP from IPv4 layer if not bound
596                let ip = get_local_ip_bytes();
597                // Allocate ephemeral port for unbound socket
598                let port = self.allocate_port();
599                self.register_port(port, socket.self_weak.clone());
600                (ip, port)
601            }
602        };
603
604        let total_length = (8 + data.len()) as u16;
605
606        let mut header = UdpHeader::new(src_port, dest_port, total_length);
607
608        header.checksum = header.calculate_checksum(src_ip_bytes, dest_ip, &data);
609
610        let mut udp_packet = Vec::with_capacity(8 + data.len());
611        udp_packet.extend_from_slice(&header.to_bytes());
612        udp_packet.extend_from_slice(&data);
613
614        let mut ip_context = LayerContext::new();
615        ip_context.set("ip_src", &src_ip_bytes);
616        ip_context.set("ip_dst", &dest_ip);
617        ip_context.set("ip_protocol", &[17]);
618
619        early_println!(
620            "[UDP] Send: {} bytes (src port: {}, dst: {}.{}.{}.{})",
621            udp_packet.len(),
622            src_port,
623            dest_ip[0],
624            dest_ip[1],
625            dest_ip[2],
626            dest_ip[3]
627        );
628
629        if let Some(ip_layer) = get_network_manager().get_layer("ip") {
630            ip_layer.send(&udp_packet, &ip_context, &[])?;
631        }
632
633        let mut stats = self.stats.write();
634        stats.packets_sent += 1;
635        stats.bytes_sent += udp_packet.len() as u64;
636
637        Ok(())
638    }
639
640    /// Receive a UDP datagram
641    pub fn receive_datagram(
642        &self,
643        src_ip: Ipv4Address,
644        src_port: u16,
645        dst_port: u16,
646        data: Vec<u8>,
647    ) {
648        let mut stats = self.stats.write();
649        stats.packets_received += 1;
650        stats.bytes_received += (8 + data.len()) as u64;
651
652        if let Some(socket) = self.find_socket(dst_port) {
653            socket.deliver_datagram(data);
654            let mut remote_lock = socket.remote_addr.write();
655            if remote_lock.is_none() {
656                *remote_lock = Some(SocketAddress::Inet(Inet4SocketAddress::new(
657                    src_ip.0, src_port,
658                )));
659            }
660        }
661    }
662}
663
664impl NetworkLayer for UdpLayer {
665    fn register_protocol(&self, _proto_num: u16, _handler: Arc<dyn NetworkLayer>) {
666        // UDP is typically a leaf protocol
667    }
668
669    fn send(
670        &self,
671        _packet: &[u8],
672        _context: &LayerContext,
673        _next_layers: &[Arc<dyn NetworkLayer>],
674    ) -> Result<(), SocketError> {
675        // UDP send is handled through send_datagram method
676        Ok(())
677    }
678
679    fn receive(&self, packet: &[u8], context: Option<&LayerContext>) -> Result<(), SocketError> {
680        let mut src_ip = Ipv4Address::new(0, 0, 0, 0);
681        let mut dst_ip = Ipv4Address::new(0, 0, 0, 0);
682        if let Some(ctx) = context {
683            if let Some(raw) = ctx.get("ip_src") {
684                if raw.len() >= 4 {
685                    src_ip = Ipv4Address::new(raw[0], raw[1], raw[2], raw[3]);
686                }
687            }
688            if let Some(raw) = ctx.get("ip_dst") {
689                if raw.len() >= 4 {
690                    dst_ip = Ipv4Address::new(raw[0], raw[1], raw[2], raw[3]);
691                }
692            }
693        }
694        self.receive_packet(src_ip, dst_ip, packet)
695    }
696
697    fn name(&self) -> &'static str {
698        "UDP"
699    }
700
701    fn stats(&self) -> NetworkLayerStats {
702        self.stats.read().clone()
703    }
704
705    fn as_any(&self) -> &dyn core::any::Any {
706        self
707    }
708}
709
710impl UdpLayer {
711    /// Receive a UDP datagram
712    pub fn receive_packet(
713        &self,
714        src_ip: Ipv4Address,
715        _dst_ip: Ipv4Address,
716        packet: &[u8],
717    ) -> Result<(), SocketError> {
718        if packet.len() < 8 {
719            return Err(SocketError::InvalidPacket);
720        }
721
722        // Parse UDP header
723        let header = UdpHeader::from_bytes(&packet[..8]).ok_or(SocketError::InvalidPacket)?;
724
725        let data_offset = header.length as usize;
726        if data_offset < 8 || data_offset > packet.len() {
727            return Err(SocketError::InvalidPacket);
728        }
729
730        let data = &packet[8..data_offset];
731
732        // Receive the datagram
733        self.receive_datagram(src_ip, header.src_port, header.dst_port, data.to_vec());
734
735        Ok(())
736    }
737}
738
739#[cfg(test)]
740mod tests {
741    use super::*;
742
743    #[test_case]
744    fn test_udp_header_creation() {
745        let header = UdpHeader::new(1234, 5678, 100);
746
747        let src_port = unsafe { core::ptr::addr_of!(header.src_port).read_unaligned() };
748        let dst_port = unsafe { core::ptr::addr_of!(header.dst_port).read_unaligned() };
749        let length = unsafe { core::ptr::addr_of!(header.length).read_unaligned() };
750        assert_eq!(src_port, 1234);
751        assert_eq!(dst_port, 5678);
752        assert_eq!(length, 100);
753    }
754
755    #[test_case]
756    fn test_udp_header_serialization() {
757        let header = UdpHeader::new(1234, 5678, 100);
758        let bytes = header.to_bytes();
759
760        assert_eq!(bytes.len(), 8);
761        assert_eq!(u16::from_be_bytes([bytes[0], bytes[1]]), 1234);
762        assert_eq!(u16::from_be_bytes([bytes[2], bytes[3]]), 5678);
763        assert_eq!(u16::from_be_bytes([bytes[4], bytes[5]]), 100);
764    }
765
766    #[test_case]
767    fn test_udp_header_parsing() {
768        let mut bytes = [0u8; 8];
769        bytes[0..2].copy_from_slice(&1234u16.to_be_bytes());
770        bytes[2..4].copy_from_slice(&5678u16.to_be_bytes());
771        bytes[4..6].copy_from_slice(&100u16.to_be_bytes());
772        bytes[6..8].copy_from_slice(&0xABCDu16.to_be_bytes());
773
774        let header = UdpHeader::from_bytes(&bytes).unwrap();
775
776        let src_port = unsafe { core::ptr::addr_of!(header.src_port).read_unaligned() };
777        let dst_port = unsafe { core::ptr::addr_of!(header.dst_port).read_unaligned() };
778        let length = unsafe { core::ptr::addr_of!(header.length).read_unaligned() };
779        let checksum = unsafe { core::ptr::addr_of!(header.checksum).read_unaligned() };
780        assert_eq!(src_port, 1234);
781        assert_eq!(dst_port, 5678);
782        assert_eq!(length, 100);
783        assert_eq!(checksum, 0xABCD);
784    }
785
786    #[test_case]
787    fn test_udp_checksum() {
788        let src_ip = [192, 168, 1, 100];
789        let dst_ip = [192, 168, 1, 1];
790        let data = b"test";
791
792        let mut header = UdpHeader::new(1234, 5678, (8 + data.len()) as u16);
793        header.checksum = header.calculate_checksum(src_ip, dst_ip, data);
794
795        // Just verify that checksum calculation runs without panicking
796        let checksum = unsafe { core::ptr::addr_of!(header.checksum).read_unaligned() };
797        assert_ne!(checksum, 0);
798    }
799
800    #[test_case]
801    fn test_udp_layer_creation() {
802        let udp_layer = UdpLayer::new();
803
804        // Test port allocation
805        let port1 = udp_layer.allocate_port();
806        let port2 = udp_layer.allocate_port();
807
808        assert!(port1 >= 49152 && port1 <= 65535);
809        assert!(port2 >= 49152 && port2 <= 65535);
810        assert_ne!(port1, port2);
811    }
812}
kernel/network/udp.rs

kernel/network/
udp.rs
