1 files changed, 17 insertions, 8 deletions

diff --git a/tests/libqos/virtio.c b/tests/libqos/virtio.c
index 0879a621c8..0dad5c19ac 100644
--- a/tests/libqos/virtio.c
+++ b/tests/libqos/virtio.c
@@ -119,6 +119,8 @@ uint8_t qvirtio_wait_status_byte_no_isr(QVirtioDevice *d,
 /*
  * qvirtio_wait_used_elem:
  * @desc_idx: The next expected vq->desc[] index in the used ring
+ * @len: A pointer that is filled with the length written into the buffer, may
+ *       be NULL
  * @timeout_us: How many microseconds to wait before failing
  *
  * This function waits for the next completed request on the used ring.
@@ -126,6 +128,7 @@ uint8_t qvirtio_wait_status_byte_no_isr(QVirtioDevice *d,
 void qvirtio_wait_used_elem(QVirtioDevice *d,
                             QVirtQueue *vq,
                             uint32_t desc_idx,
+                            uint32_t *len,
                             gint64 timeout_us)
 {
     gint64 start_time = g_get_monotonic_time();
@@ -136,7 +139,7 @@ void qvirtio_wait_used_elem(QVirtioDevice *d,
         clock_step(100);
 
         if (d->bus->get_queue_isr_status(d, vq) &&
-            qvirtqueue_get_buf(vq, &got_desc_idx)) {
+            qvirtqueue_get_buf(vq, &got_desc_idx, len)) {
             g_assert_cmpint(got_desc_idx, ==, desc_idx);
             return;
         }
@@ -304,30 +307,36 @@ void qvirtqueue_kick(QVirtioDevice *d, QVirtQueue *vq, uint32_t free_head)
 /*
  * qvirtqueue_get_buf:
  * @desc_idx: A pointer that is filled with the vq->desc[] index, may be NULL
+ * @len: A pointer that is filled with the length written into the buffer, may
+ *       be NULL
  *
  * This function gets the next used element if there is one ready.
  *
  * Returns: true if an element was ready, false otherwise
  */
-bool qvirtqueue_get_buf(QVirtQueue *vq, uint32_t *desc_idx)
+bool qvirtqueue_get_buf(QVirtQueue *vq, uint32_t *desc_idx, uint32_t *len)
 {
     uint16_t idx;
+    uint64_t elem_addr;
 
     idx = readw(vq->used + offsetof(struct vring_used, idx));
     if (idx == vq->last_used_idx) {
         return false;
     }
 
-    if (desc_idx) {
-        uint64_t elem_addr;
+    elem_addr = vq->used +
+        offsetof(struct vring_used, ring) +
+        (vq->last_used_idx % vq->size) *
+        sizeof(struct vring_used_elem);
 
-        elem_addr = vq->used +
-                    offsetof(struct vring_used, ring) +
-                    (vq->last_used_idx % vq->size) *
-                    sizeof(struct vring_used_elem);
+    if (desc_idx) {
         *desc_idx = readl(elem_addr + offsetof(struct vring_used_elem, id));
     }
 
+    if (len) {
+        *len = readw(elem_addr + offsetof(struct vring_used_elem, len));
+    }
+
     vq->last_used_idx++;
     return true;
 }