Browse Source

Add blocking helper methods.

Andrew Walbran 2 years ago
parent
commit
277783ed87
2 changed files with 28 additions and 1 deletions
  1. 2 1
      examples/aarch64/src/main.rs
  2. 26 0
      src/device/socket/vsock.rs

+ 2 - 1
examples/aarch64/src/main.rs

@@ -194,13 +194,14 @@ fn virtio_socket<T: Transport>(transport: T) -> virtio_drivers::Result<()> {
     let port = 1221;
     info!("Connecting to host on port {port}...");
     socket.connect(host_cid, port, port)?;
+    socket.wait_for_connect()?;
     info!("Connected to the host");
 
     const EXCHANGE_NUM: usize = 2;
     let messages = ["0-Ack. Hello from guest.", "1-Ack. Received again."];
     for k in 0..EXCHANGE_NUM {
         let mut buffer = [0u8; 24];
-        let socket_event = socket.poll_recv(&mut buffer)?;
+        let socket_event = socket.wait_for_recv(&mut buffer)?;
         let VsockEventType::Received {length, ..} = socket_event.event_type else {
             panic!("Received unexpected socket event {:?}", socket_event);
         };

+ 26 - 0
src/device/socket/vsock.rs

@@ -9,6 +9,7 @@ use crate::queue::VirtQueue;
 use crate::transport::Transport;
 use crate::volatile::volread;
 use crate::Result;
+use core::hint::spin_loop;
 use core::mem::size_of;
 use core::ptr::NonNull;
 use log::{debug, info};
@@ -224,6 +225,17 @@ impl<H: Hal, T: Transport> VirtIOSocket<H, T> {
         Ok(())
     }
 
+    /// Blocks until the peer either accepts our connection request (with a
+    /// `VIRTIO_VSOCK_OP_RESPONSE`) or rejects it (with a
+    /// `VIRTIO_VSOCK_OP_RST`).
+    pub fn wait_for_connect(&mut self) -> Result {
+        match self.wait_for_recv(&mut [])?.event_type {
+            VsockEventType::Connected => Ok(()),
+            VsockEventType::Disconnected { .. } => Err(SocketError::ConnectionFailed.into()),
+            VsockEventType::Received { .. } => Err(SocketError::InvalidOperation.into()),
+        }
+    }
+
     /// Requests the peer to send us a credit update for the current connection.
     fn request_credit(&mut self) -> Result {
         let connection_info = self.connection_info()?;
@@ -290,6 +302,20 @@ impl<H: Hal, T: Transport> VirtIOSocket<H, T> {
         Ok(event)
     }
 
+    /// Blocks until we get some event from the vsock device.
+    ///
+    /// A buffer must be provided to put the data in if there is some to
+    /// receive.
+    pub fn wait_for_recv(&mut self, buffer: &mut [u8]) -> Result<VsockEvent> {
+        loop {
+            if let Some(event) = self.poll_recv(buffer)? {
+                return Ok(event);
+            } else {
+                spin_loop();
+            }
+        }
+    }
+
     /// Request to shut down the connection cleanly.
     ///
     /// This returns as soon as the request is sent; you should wait until `poll_recv` returns a