Skip to main content

arcbox_virtio_vsock/
connection.rs

1//! Vsock connection state — used by the in-process loopback path.
2
3use crate::addr::VsockAddr;
4
5/// Connection state.
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum ConnectionState {
8    /// Initial state.
9    Idle,
10    /// Connection requested.
11    Connecting,
12    /// Connected.
13    Connected,
14    /// Shutting down.
15    Closing,
16    /// Closed.
17    Closed,
18}
19
20/// A vsock connection.
21#[derive(Debug)]
22#[allow(dead_code)]
23pub struct VsockConnection {
24    /// Local address.
25    pub local: VsockAddr,
26    /// Remote address.
27    pub remote: VsockAddr,
28    /// Connection state.
29    pub state: ConnectionState,
30    /// Receive buffer.
31    rx_buf: Vec<u8>,
32    /// Transmit buffer.
33    tx_buf: Vec<u8>,
34    /// Buffer allocation (credit).
35    pub(crate) buf_alloc: u32,
36    /// Forward count.
37    pub(crate) fwd_cnt: u32,
38    /// Peer buffer allocation.
39    pub(crate) peer_buf_alloc: u32,
40    /// Peer forward count.
41    pub(crate) peer_fwd_cnt: u32,
42}
43
44impl VsockConnection {
45    /// Creates a new connection.
46    #[must_use]
47    pub fn new(local: VsockAddr, remote: VsockAddr) -> Self {
48        Self {
49            local,
50            remote,
51            state: ConnectionState::Idle,
52            rx_buf: Vec::with_capacity(64 * 1024),
53            tx_buf: Vec::with_capacity(64 * 1024),
54            buf_alloc: 64 * 1024,
55            fwd_cnt: 0,
56            peer_buf_alloc: 0,
57            peer_fwd_cnt: 0,
58        }
59    }
60
61    /// Returns bytes available to send.
62    #[must_use]
63    pub fn tx_available(&self) -> usize {
64        self.tx_buf.len()
65    }
66
67    /// Returns bytes available to receive.
68    #[must_use]
69    pub fn rx_available(&self) -> usize {
70        self.rx_buf.len()
71    }
72
73    /// Enqueues data for transmission.
74    pub fn enqueue_tx(&mut self, data: &[u8]) {
75        self.tx_buf.extend_from_slice(data);
76    }
77
78    /// Dequeues transmitted data.
79    pub fn dequeue_tx(&mut self, max_len: usize) -> Vec<u8> {
80        let len = max_len.min(self.tx_buf.len());
81        self.tx_buf.drain(..len).collect()
82    }
83
84    /// Enqueues received data.
85    pub fn enqueue_rx(&mut self, data: &[u8]) {
86        self.rx_buf.extend_from_slice(data);
87        self.fwd_cnt = self.fwd_cnt.wrapping_add(data.len() as u32);
88    }
89
90    /// Dequeues received data.
91    pub fn dequeue_rx(&mut self, max_len: usize) -> Vec<u8> {
92        let len = max_len.min(self.rx_buf.len());
93        self.rx_buf.drain(..len).collect()
94    }
95
96    /// Updates peer credit info.
97    pub const fn update_peer_credit(&mut self, buf_alloc: u32, fwd_cnt: u32) {
98        self.peer_buf_alloc = buf_alloc;
99        self.peer_fwd_cnt = fwd_cnt;
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[test]
108    fn test_vsock_connection_state() {
109        let conn = VsockConnection::new(VsockAddr::new(3, 1000), VsockAddr::new(2, 80));
110        assert_eq!(conn.state, ConnectionState::Idle);
111        assert_eq!(conn.tx_available(), 0);
112        assert_eq!(conn.rx_available(), 0);
113    }
114
115    #[test]
116    fn test_vsock_connection_buffers() {
117        let mut conn = VsockConnection::new(VsockAddr::new(3, 1000), VsockAddr::new(2, 80));
118
119        conn.enqueue_tx(b"hello");
120        assert_eq!(conn.tx_available(), 5);
121        let data = conn.dequeue_tx(3);
122        assert_eq!(&data, b"hel");
123        assert_eq!(conn.tx_available(), 2);
124
125        conn.enqueue_rx(b"world");
126        assert_eq!(conn.rx_available(), 5);
127        let data = conn.dequeue_rx(10);
128        assert_eq!(&data, b"world");
129        assert_eq!(conn.rx_available(), 0);
130    }
131}