arcbox_virtio_vsock/
connection.rs1use crate::addr::VsockAddr;
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum ConnectionState {
8 Idle,
10 Connecting,
12 Connected,
14 Closing,
16 Closed,
18}
19
20#[derive(Debug)]
22#[allow(dead_code)]
23pub struct VsockConnection {
24 pub local: VsockAddr,
26 pub remote: VsockAddr,
28 pub state: ConnectionState,
30 rx_buf: Vec<u8>,
32 tx_buf: Vec<u8>,
34 pub(crate) buf_alloc: u32,
36 pub(crate) fwd_cnt: u32,
38 pub(crate) peer_buf_alloc: u32,
40 pub(crate) peer_fwd_cnt: u32,
42}
43
44impl VsockConnection {
45 #[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 #[must_use]
63 pub fn tx_available(&self) -> usize {
64 self.tx_buf.len()
65 }
66
67 #[must_use]
69 pub fn rx_available(&self) -> usize {
70 self.rx_buf.len()
71 }
72
73 pub fn enqueue_tx(&mut self, data: &[u8]) {
75 self.tx_buf.extend_from_slice(data);
76 }
77
78 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 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 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 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}