amadeus_node/
socket.rs

1use crate::metrics::Metrics;
2use std::io::Result;
3use std::net::SocketAddr;
4
5/// Trait for UDP socket operations with optional dump/replay functionality
6#[async_trait::async_trait]
7pub trait UdpSocketExt: Send + Sync {
8    async fn recv_from(&self, buf: &mut [u8]) -> Result<(usize, SocketAddr)>;
9    async fn send_to(&self, buf: &[u8], target: SocketAddr) -> Result<usize>;
10
11    /// Wrapper function that calls recv_from and increments metrics
12    async fn recv_from_with_metrics(&self, buf: &mut [u8], metrics: &Metrics) -> Result<(usize, SocketAddr)> {
13        let result = self.recv_from(buf).await?;
14        metrics.add_incoming_udp_packet(result.0);
15        Ok(result)
16    }
17
18    /// Wrapper function that calls send_to and increments metrics
19    async fn send_to_with_metrics(&self, buf: &[u8], target: SocketAddr, metrics: &Metrics) -> Result<usize> {
20        let result = self.send_to(buf, target).await?;
21        metrics.add_outgoing_udp_packet(result);
22        Ok(result)
23    }
24}
25
26#[cfg(test)]
27pub struct MockSocket {
28    recv_data: Vec<u8>,
29    recv_addr: SocketAddr,
30}
31
32#[cfg(test)]
33impl MockSocket {
34    pub fn new() -> Self {
35        Self::with_data_and_src(vec![], "127.0.0.1:1234".parse().unwrap())
36    }
37
38    fn with_data_and_src(recv_data: Vec<u8>, recv_addr: SocketAddr) -> Self {
39        Self { recv_data, recv_addr }
40    }
41}
42
43#[cfg(test)]
44#[async_trait::async_trait]
45impl UdpSocketExt for MockSocket {
46    async fn recv_from(&self, buf: &mut [u8]) -> Result<(usize, SocketAddr)> {
47        let len = self.recv_data.len().min(buf.len());
48        buf[..len].copy_from_slice(&self.recv_data[..len]);
49        Ok((len, self.recv_addr))
50    }
51
52    async fn send_to(&self, buf: &[u8], _target: SocketAddr) -> Result<usize> {
53        Ok(buf.len())
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60    use crate::metrics::Metrics;
61
62    #[tokio::test]
63    async fn test_recv_from_with_metrics() {
64        let socket = MockSocket::with_data_and_src(vec![1, 2, 3, 4], "127.0.0.1:1234".parse().unwrap());
65        let metrics = Metrics::new();
66        let mut buf = vec![0u8; 10];
67
68        let (len, addr) = socket.recv_from_with_metrics(&mut buf, &metrics).await.unwrap();
69
70        assert_eq!(len, 4);
71        assert_eq!(addr.to_string(), "127.0.0.1:1234");
72
73        // Check that metrics were incremented
74        let snapshot = metrics.get_snapshot();
75        assert_eq!(snapshot.udp.incoming_packets, 1);
76        assert_eq!(snapshot.udp.incoming_bytes, 4);
77    }
78
79    #[tokio::test]
80    async fn test_send_to_with_metrics() {
81        let socket = MockSocket::with_data_and_src(vec![1, 2, 3, 4], "127.0.0.1:1234".parse().unwrap());
82        let metrics = Metrics::new();
83        let data = vec![1, 2, 3, 4, 5];
84        let target = "127.0.0.1:5678".parse().unwrap();
85
86        let sent = socket.send_to_with_metrics(&data, target, &metrics).await.unwrap();
87
88        assert_eq!(sent, 5);
89
90        // Check that metrics were incremented
91        let snapshot = metrics.get_snapshot();
92        assert_eq!(snapshot.udp.outgoing_packets, 1);
93        assert_eq!(snapshot.udp.outgoing_bytes, 5);
94    }
95}
amadeus_node/socket.rs

amadeus_node/
socket.rs
