multiconnectionmanager.rs 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519
  1. use super::{
  2. protocol::VsockAddr, vsock::ConnectionInfo, SocketError, VirtIOSocket, VsockEvent,
  3. VsockEventType,
  4. };
  5. use crate::{transport::Transport, Hal, Result};
  6. use alloc::{boxed::Box, vec::Vec};
  7. use core::cmp::min;
  8. use core::convert::TryInto;
  9. use core::hint::spin_loop;
  10. use log::debug;
  11. use zerocopy::FromBytes;
  12. const PER_CONNECTION_BUFFER_CAPACITY: usize = 1024;
  13. /// A higher level interface for VirtIO socket (vsock) devices.
  14. ///
  15. /// This keeps track of multiple vsock connections.
  16. ///
  17. /// # Example
  18. ///
  19. /// ```
  20. /// # use virtio_drivers::{Error, Hal};
  21. /// # use virtio_drivers::transport::Transport;
  22. /// use virtio_drivers::device::socket::{VirtIOSocket, VsockAddr, VsockConnectionManager};
  23. ///
  24. /// # fn example<HalImpl: Hal, T: Transport>(transport: T) -> Result<(), Error> {
  25. /// let mut socket = VsockConnectionManager::new(VirtIOSocket::<HalImpl, _>::new(transport)?);
  26. ///
  27. /// // Start a thread to call `socket.poll()` and handle events.
  28. ///
  29. /// let remote_address = VsockAddr { cid: 2, port: 42 };
  30. /// let local_port = 1234;
  31. /// socket.connect(remote_address, local_port)?;
  32. ///
  33. /// // Wait until `socket.poll()` returns an event indicating that the socket is connected.
  34. ///
  35. /// socket.send(remote_address, local_port, "Hello world".as_bytes())?;
  36. ///
  37. /// socket.shutdown(remote_address, local_port)?;
  38. /// # Ok(())
  39. /// # }
  40. /// ```
  41. pub struct VsockConnectionManager<H: Hal, T: Transport> {
  42. driver: VirtIOSocket<H, T>,
  43. connections: Vec<Connection>,
  44. }
  45. #[derive(Debug)]
  46. struct Connection {
  47. info: ConnectionInfo,
  48. buffer: RingBuffer,
  49. }
  50. impl<H: Hal, T: Transport> VsockConnectionManager<H, T> {
  51. /// Construct a new connection manager wrapping the given low-level VirtIO socket driver.
  52. pub fn new(driver: VirtIOSocket<H, T>) -> Self {
  53. Self {
  54. driver,
  55. connections: Vec::new(),
  56. }
  57. }
  58. /// Returns the CID which has been assigned to this guest.
  59. pub fn guest_cid(&self) -> u64 {
  60. self.driver.guest_cid()
  61. }
  62. /// Sends a request to connect to the given destination.
  63. ///
  64. /// This returns as soon as the request is sent; you should wait until `poll` returns a
  65. /// `VsockEventType::Connected` event indicating that the peer has accepted the connection
  66. /// before sending data.
  67. pub fn connect(&mut self, destination: VsockAddr, src_port: u32) -> Result {
  68. if self.connections.iter().any(|connection| {
  69. connection.info.dst == destination && connection.info.src_port == src_port
  70. }) {
  71. return Err(SocketError::ConnectionExists.into());
  72. }
  73. let mut new_connection_info = ConnectionInfo::new(destination, src_port);
  74. new_connection_info.buf_alloc = PER_CONNECTION_BUFFER_CAPACITY.try_into().unwrap();
  75. self.driver.connect(&new_connection_info)?;
  76. debug!("Connection requested: {:?}", new_connection_info);
  77. self.connections.push(Connection {
  78. info: new_connection_info,
  79. buffer: RingBuffer::new(PER_CONNECTION_BUFFER_CAPACITY),
  80. });
  81. Ok(())
  82. }
  83. /// Sends the buffer to the destination.
  84. pub fn send(&mut self, destination: VsockAddr, src_port: u32, buffer: &[u8]) -> Result {
  85. let connection = self
  86. .connections
  87. .iter_mut()
  88. .find(|connection| {
  89. connection.info.dst == destination && connection.info.src_port == src_port
  90. })
  91. .ok_or(SocketError::NotConnected)?;
  92. self.driver.send(buffer, &mut connection.info)
  93. }
  94. /// Polls the vsock device to receive data or other updates.
  95. pub fn poll(&mut self) -> Result<Option<VsockEvent>> {
  96. let guest_cid = self.driver.guest_cid();
  97. let connections = &mut self.connections;
  98. self.driver.poll(|event, body| {
  99. let connection = connections
  100. .iter_mut()
  101. .find(|connection| event.matches_connection(&connection.info, guest_cid));
  102. let Some(connection) = connection else {
  103. // Skip events which don't match any connection we know about.
  104. return Ok(None);
  105. };
  106. // Update stored connection info.
  107. connection.info.update_for_event(&event);
  108. match event.event_type {
  109. VsockEventType::ConnectionRequest => {
  110. // TODO: Send Rst or handle incoming connections.
  111. }
  112. VsockEventType::Connected => {}
  113. VsockEventType::Disconnected { .. } => {
  114. // TODO: Wait until client reads all data before removing connection.
  115. //self.connection_info = None;
  116. }
  117. VsockEventType::Received { length } => {
  118. // Copy to buffer
  119. if !connection.buffer.write(body) {
  120. return Err(SocketError::OutputBufferTooShort(length).into());
  121. }
  122. }
  123. VsockEventType::CreditRequest => {
  124. // TODO: Send a credit update.
  125. }
  126. VsockEventType::CreditUpdate => {}
  127. }
  128. Ok(Some(event))
  129. })
  130. }
  131. /// Reads data received from the given connection.
  132. pub fn recv(&mut self, peer: VsockAddr, src_port: u32, buffer: &mut [u8]) -> Result<usize> {
  133. let connection = self
  134. .connections
  135. .iter_mut()
  136. .find(|connection| connection.info.dst == peer && connection.info.src_port == src_port)
  137. .ok_or(SocketError::NotConnected)?;
  138. // Copy from ring buffer
  139. let bytes_read = connection.buffer.read(buffer);
  140. connection.info.done_forwarding(bytes_read);
  141. Ok(bytes_read)
  142. }
  143. /// Blocks until we get some event from the vsock device.
  144. pub fn wait_for_event(&mut self) -> Result<VsockEvent> {
  145. loop {
  146. if let Some(event) = self.poll()? {
  147. return Ok(event);
  148. } else {
  149. spin_loop();
  150. }
  151. }
  152. }
  153. /// Requests to shut down the connection cleanly.
  154. ///
  155. /// This returns as soon as the request is sent; you should wait until `poll` returns a
  156. /// `VsockEventType::Disconnected` event if you want to know that the peer has acknowledged the
  157. /// shutdown.
  158. pub fn shutdown(&mut self, destination: VsockAddr, src_port: u32) -> Result {
  159. let connection = self
  160. .connections
  161. .iter()
  162. .find(|connection| {
  163. connection.info.dst == destination && connection.info.src_port == src_port
  164. })
  165. .ok_or(SocketError::NotConnected)?;
  166. self.driver.shutdown(&connection.info)
  167. }
  168. /// Forcibly closes the connection without waiting for the peer.
  169. pub fn force_close(&mut self, destination: VsockAddr, src_port: u32) -> Result {
  170. let (index, connection) = self
  171. .connections
  172. .iter()
  173. .enumerate()
  174. .find(|(_, connection)| {
  175. connection.info.dst == destination && connection.info.src_port == src_port
  176. })
  177. .ok_or(SocketError::NotConnected)?;
  178. self.driver.force_close(&connection.info)?;
  179. self.connections.swap_remove(index);
  180. Ok(())
  181. }
  182. }
  183. #[derive(Debug)]
  184. struct RingBuffer {
  185. buffer: Box<[u8]>,
  186. /// The number of bytes currently in the buffer.
  187. used: usize,
  188. /// The index of the first used byte in the buffer.
  189. start: usize,
  190. }
  191. impl RingBuffer {
  192. pub fn new(capacity: usize) -> Self {
  193. Self {
  194. buffer: FromBytes::new_box_slice_zeroed(capacity),
  195. used: 0,
  196. start: 0,
  197. }
  198. }
  199. /// Returns the number of bytes currently used in the buffer.
  200. pub fn used(&self) -> usize {
  201. self.used
  202. }
  203. /// Returns the number of bytes currently free in the buffer.
  204. pub fn available(&self) -> usize {
  205. self.buffer.len() - self.used
  206. }
  207. /// Adds the given bytes to the buffer if there is enough capacity for them all.
  208. ///
  209. /// Returns true if they were added, or false if they were not.
  210. pub fn write(&mut self, bytes: &[u8]) -> bool {
  211. if bytes.len() > self.available() {
  212. return false;
  213. }
  214. let end = (self.start + self.used) % self.buffer.len();
  215. let write_before_wraparound = min(bytes.len(), self.buffer.len() - end);
  216. let write_after_wraparound = bytes
  217. .len()
  218. .checked_sub(write_before_wraparound)
  219. .unwrap_or_default();
  220. self.buffer[end..end + write_before_wraparound]
  221. .copy_from_slice(&bytes[0..write_before_wraparound]);
  222. self.buffer[0..write_after_wraparound].copy_from_slice(&bytes[write_before_wraparound..]);
  223. self.used += bytes.len();
  224. true
  225. }
  226. /// Reads and removes as many bytes as possible from the buffer, up to the length of the given
  227. /// buffer.
  228. pub fn read(&mut self, out: &mut [u8]) -> usize {
  229. let bytes_read = min(self.used, out.len());
  230. // The number of bytes to copy out between `start` and the end of the buffer.
  231. let read_before_wraparound = min(bytes_read, self.buffer.len() - self.start);
  232. // The number of bytes to copy out from the beginning of the buffer after wrapping around.
  233. let read_after_wraparound = bytes_read
  234. .checked_sub(read_before_wraparound)
  235. .unwrap_or_default();
  236. out[0..read_before_wraparound]
  237. .copy_from_slice(&self.buffer[self.start..self.start + read_before_wraparound]);
  238. out[read_before_wraparound..bytes_read]
  239. .copy_from_slice(&self.buffer[0..read_after_wraparound]);
  240. self.used -= bytes_read;
  241. self.start = (self.start + bytes_read) % self.buffer.len();
  242. bytes_read
  243. }
  244. }
  245. #[cfg(test)]
  246. mod tests {
  247. use super::*;
  248. use crate::{
  249. device::socket::{
  250. protocol::{SocketType, VirtioVsockConfig, VirtioVsockHdr, VirtioVsockOp},
  251. vsock::{VsockBufferStatus, QUEUE_SIZE, RX_QUEUE_IDX, TX_QUEUE_IDX},
  252. },
  253. hal::fake::FakeHal,
  254. transport::{
  255. fake::{FakeTransport, QueueStatus, State},
  256. DeviceStatus, DeviceType,
  257. },
  258. volatile::ReadOnly,
  259. };
  260. use alloc::{sync::Arc, vec};
  261. use core::{mem::size_of, ptr::NonNull};
  262. use std::{sync::Mutex, thread};
  263. use zerocopy::{AsBytes, FromBytes};
  264. #[test]
  265. fn send_recv() {
  266. let host_cid = 2;
  267. let guest_cid = 66;
  268. let host_port = 1234;
  269. let guest_port = 4321;
  270. let host_address = VsockAddr {
  271. cid: host_cid,
  272. port: host_port,
  273. };
  274. let hello_from_guest = "Hello from guest";
  275. let hello_from_host = "Hello from host";
  276. let mut config_space = VirtioVsockConfig {
  277. guest_cid_low: ReadOnly::new(66),
  278. guest_cid_high: ReadOnly::new(0),
  279. };
  280. let state = Arc::new(Mutex::new(State {
  281. status: DeviceStatus::empty(),
  282. driver_features: 0,
  283. guest_page_size: 0,
  284. interrupt_pending: false,
  285. queues: vec![
  286. QueueStatus::default(),
  287. QueueStatus::default(),
  288. QueueStatus::default(),
  289. ],
  290. }));
  291. let transport = FakeTransport {
  292. device_type: DeviceType::Socket,
  293. max_queue_size: 32,
  294. device_features: 0,
  295. config_space: NonNull::from(&mut config_space),
  296. state: state.clone(),
  297. };
  298. let mut socket = VsockConnectionManager::new(
  299. VirtIOSocket::<FakeHal, FakeTransport<VirtioVsockConfig>>::new(transport).unwrap(),
  300. );
  301. // Start a thread to simulate the device.
  302. let handle = thread::spawn(move || {
  303. // Wait for connection request.
  304. State::wait_until_queue_notified(&state, TX_QUEUE_IDX);
  305. assert_eq!(
  306. VirtioVsockHdr::read_from(
  307. state
  308. .lock()
  309. .unwrap()
  310. .read_from_queue::<QUEUE_SIZE>(TX_QUEUE_IDX)
  311. .as_slice()
  312. )
  313. .unwrap(),
  314. VirtioVsockHdr {
  315. op: VirtioVsockOp::Request.into(),
  316. src_cid: guest_cid.into(),
  317. dst_cid: host_cid.into(),
  318. src_port: guest_port.into(),
  319. dst_port: host_port.into(),
  320. len: 0.into(),
  321. socket_type: SocketType::Stream.into(),
  322. flags: 0.into(),
  323. buf_alloc: 1024.into(),
  324. fwd_cnt: 0.into(),
  325. }
  326. );
  327. // Accept connection and give the peer enough credit to send the message.
  328. state.lock().unwrap().write_to_queue::<QUEUE_SIZE>(
  329. RX_QUEUE_IDX,
  330. VirtioVsockHdr {
  331. op: VirtioVsockOp::Response.into(),
  332. src_cid: host_cid.into(),
  333. dst_cid: guest_cid.into(),
  334. src_port: host_port.into(),
  335. dst_port: guest_port.into(),
  336. len: 0.into(),
  337. socket_type: SocketType::Stream.into(),
  338. flags: 0.into(),
  339. buf_alloc: 50.into(),
  340. fwd_cnt: 0.into(),
  341. }
  342. .as_bytes(),
  343. );
  344. // Expect the guest to send some data.
  345. State::wait_until_queue_notified(&state, TX_QUEUE_IDX);
  346. let request = state
  347. .lock()
  348. .unwrap()
  349. .read_from_queue::<QUEUE_SIZE>(TX_QUEUE_IDX);
  350. assert_eq!(
  351. request.len(),
  352. size_of::<VirtioVsockHdr>() + hello_from_guest.len()
  353. );
  354. assert_eq!(
  355. VirtioVsockHdr::read_from_prefix(request.as_slice()).unwrap(),
  356. VirtioVsockHdr {
  357. op: VirtioVsockOp::Rw.into(),
  358. src_cid: guest_cid.into(),
  359. dst_cid: host_cid.into(),
  360. src_port: guest_port.into(),
  361. dst_port: host_port.into(),
  362. len: (hello_from_guest.len() as u32).into(),
  363. socket_type: SocketType::Stream.into(),
  364. flags: 0.into(),
  365. buf_alloc: 1024.into(),
  366. fwd_cnt: 0.into(),
  367. }
  368. );
  369. assert_eq!(
  370. &request[size_of::<VirtioVsockHdr>()..],
  371. hello_from_guest.as_bytes()
  372. );
  373. println!("Host sending");
  374. // Send a response.
  375. let mut response = vec![0; size_of::<VirtioVsockHdr>() + hello_from_host.len()];
  376. VirtioVsockHdr {
  377. op: VirtioVsockOp::Rw.into(),
  378. src_cid: host_cid.into(),
  379. dst_cid: guest_cid.into(),
  380. src_port: host_port.into(),
  381. dst_port: guest_port.into(),
  382. len: (hello_from_host.len() as u32).into(),
  383. socket_type: SocketType::Stream.into(),
  384. flags: 0.into(),
  385. buf_alloc: 50.into(),
  386. fwd_cnt: (hello_from_guest.len() as u32).into(),
  387. }
  388. .write_to_prefix(response.as_mut_slice());
  389. response[size_of::<VirtioVsockHdr>()..].copy_from_slice(hello_from_host.as_bytes());
  390. state
  391. .lock()
  392. .unwrap()
  393. .write_to_queue::<QUEUE_SIZE>(RX_QUEUE_IDX, &response);
  394. // Expect a shutdown.
  395. State::wait_until_queue_notified(&state, TX_QUEUE_IDX);
  396. assert_eq!(
  397. VirtioVsockHdr::read_from(
  398. state
  399. .lock()
  400. .unwrap()
  401. .read_from_queue::<QUEUE_SIZE>(TX_QUEUE_IDX)
  402. .as_slice()
  403. )
  404. .unwrap(),
  405. VirtioVsockHdr {
  406. op: VirtioVsockOp::Shutdown.into(),
  407. src_cid: guest_cid.into(),
  408. dst_cid: host_cid.into(),
  409. src_port: guest_port.into(),
  410. dst_port: host_port.into(),
  411. len: 0.into(),
  412. socket_type: SocketType::Stream.into(),
  413. flags: 0.into(),
  414. buf_alloc: 1024.into(),
  415. fwd_cnt: (hello_from_host.len() as u32).into(),
  416. }
  417. );
  418. });
  419. socket.connect(host_address, guest_port).unwrap();
  420. assert_eq!(
  421. socket.wait_for_event().unwrap(),
  422. VsockEvent {
  423. source: host_address,
  424. destination: VsockAddr {
  425. cid: guest_cid,
  426. port: guest_port,
  427. },
  428. event_type: VsockEventType::Connected,
  429. buffer_status: VsockBufferStatus {
  430. buffer_allocation: 50,
  431. forward_count: 0,
  432. },
  433. }
  434. );
  435. println!("Guest sending");
  436. socket
  437. .send(host_address, guest_port, "Hello from guest".as_bytes())
  438. .unwrap();
  439. println!("Guest waiting to receive.");
  440. assert_eq!(
  441. socket.wait_for_event().unwrap(),
  442. VsockEvent {
  443. source: host_address,
  444. destination: VsockAddr {
  445. cid: guest_cid,
  446. port: guest_port,
  447. },
  448. event_type: VsockEventType::Received {
  449. length: hello_from_host.len()
  450. },
  451. buffer_status: VsockBufferStatus {
  452. buffer_allocation: 50,
  453. forward_count: hello_from_guest.len() as u32,
  454. },
  455. }
  456. );
  457. println!("Guest getting received data.");
  458. let mut buffer = [0u8; 64];
  459. assert_eq!(
  460. socket.recv(host_address, guest_port, &mut buffer).unwrap(),
  461. hello_from_host.len()
  462. );
  463. assert_eq!(
  464. &buffer[0..hello_from_host.len()],
  465. hello_from_host.as_bytes()
  466. );
  467. socket.shutdown(host_address, guest_port).unwrap();
  468. handle.join().unwrap();
  469. }
  470. }