Take token as parameter to pop_used rather than returning it.

This makes more sense; if the device uses a different set of buffers
than the one we were expecting then we shouldn't just go ahead and pop
it, so we return an error instead.

The new peek_used method can be used to check which token is next to pop
and decide whether to handle it.

src/device/blk.rs

@@ -256,7 +256,9 @@ impl<H: Hal, T: Transport> VirtIOBlk<H, T> {
     /// ring and return it. If all completed requests have already been fetched, return
     /// Err(Error::NotReady).
     pub fn pop_used(&mut self) -> Result<u16> {
-        self.queue.pop_used().map(|p| p.0)
+        let token = self.queue.peek_used().ok_or(Error::NotReady)?;
+        self.queue.pop_used(token)?;
+        Ok(token)
     }
 
     /// Returns the size of the device's VirtQueue.

src/device/console.rs

@@ -136,8 +136,7 @@ impl<H: Hal, T: Transport> VirtIOConsole<'_, H, T> {
     fn finish_receive(&mut self) -> bool {
         let mut flag = false;
         if let Some(receive_token) = self.receive_token {
-            if let Ok((token, len)) = self.receiveq.pop_used() {
-                assert_eq!(token, receive_token);
+            if let Ok(len) = self.receiveq.pop_used(receive_token) {
                 flag = true;
                 assert_ne!(len, 0);
                 self.cursor = 0;

src/device/input.rs

@@ -64,7 +64,8 @@ impl<H: Hal, T: Transport> VirtIOInput<H, T> {
 
     /// Pop the pending event.
     pub fn pop_pending_event(&mut self) -> Option<InputEvent> {
-        if let Ok((token, _)) = self.event_queue.pop_used() {
+        if let Some(token) = self.event_queue.peek_used() {
+            self.event_queue.pop_used(token).ok()?;
             let event = &mut self.event_buf[token as usize];
             // requeue
             // Safe because buffer lasts as long as the queue.

src/lib.rs

@@ -68,6 +68,8 @@ pub enum Error {
     QueueFull,
     /// The device is not ready.
     NotReady,
+    /// The device used a different descriptor chain to the one we were expecting.
+    WrongToken,
     /// The queue is already in use.
     AlreadyUsed,
     /// Invalid parameter.

src/queue.rs

@@ -166,13 +166,12 @@ impl<H: Hal> VirtQueue<H> {
         // Notify the queue.
         transport.notify(self.queue_idx);
 
+        // Wait until there is at least one element in the used ring.
         while !self.can_pop() {
             spin_loop();
         }
-        let (popped_token, length) = self.pop_used()?;
-        assert_eq!(popped_token, token);
 
-        Ok(length)
+        self.pop_used(token)
     }
 
     /// Returns a non-null pointer to the descriptor at the given index.
@@ -191,6 +190,19 @@ impl<H: Hal> VirtQueue<H> {
         self.last_used_idx != unsafe { (*self.used.as_ptr()).idx }
     }
 
+    /// Returns the descriptor index (a.k.a. token) of the next used element without popping it, or
+    /// `None` if the used ring is empty.
+    pub fn peek_used(&self) -> Option<u16> {
+        if self.can_pop() {
+            let last_used_slot = self.last_used_idx & (self.queue_size - 1);
+            // Safe because self.used points to a valid, aligned, initialised, dereferenceable,
+            // readable instance of UsedRing.
+            Some(unsafe { (*self.used.as_ptr()).ring[last_used_slot as usize].id as u16 })
+        } else {
+            None
+        }
+    }
+
     /// Returns the number of free descriptors.
     pub fn available_desc(&self) -> usize {
         (self.queue_size - self.num_used) as usize
@@ -219,16 +231,17 @@ impl<H: Hal> VirtQueue<H> {
         }
     }
 
-    /// Get a token from device used buffers, return (token, len).
+    /// If the given token is next on the device used queue, pops it and returns the total buffer
+    /// length which was used (written) by the device.
     ///
     /// Ref: linux virtio_ring.c virtqueue_get_buf_ctx
-    pub fn pop_used(&mut self) -> Result<(u16, u32)> {
+    pub fn pop_used(&mut self, token: u16) -> Result<u32> {
         if !self.can_pop() {
             return Err(Error::NotReady);
         }
-        // read barrier
-        fence(Ordering::SeqCst);
+        // Read barrier not necessary, as can_pop already has one.
 
+        // Get the index of the start of the descriptor chain for the next element in the used ring.
         let last_used_slot = self.last_used_idx & (self.queue_size - 1);
         let index;
         let len;
@@ -239,10 +252,15 @@ impl<H: Hal> VirtQueue<H> {
             len = (*self.used.as_ptr()).ring[last_used_slot as usize].len;
         }
 
+        if index != token {
+            // The device used a different descriptor chain to the one we were expecting.
+            return Err(Error::WrongToken);
+        }
+
         self.recycle_descriptors(index);
         self.last_used_idx = self.last_used_idx.wrapping_add(1);
 
-        Ok((index, len))
+        Ok(len)
     }
 
     /// Return size of the queue.