Skip to main content

virtio/
test_helpers.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Shared test helpers for writing integration tests that drive virtio devices
5//! through descriptor rings.
6//!
7//! These helpers manipulate split virtqueue descriptor tables, available rings,
8//! and used rings in guest memory — the same operations a guest driver would
9//! perform.
10
11use crate::spec::queue::AVAIL_ELEMENT_SIZE;
12use crate::spec::queue::AVAIL_OFFSET_FLAGS;
13use crate::spec::queue::AVAIL_OFFSET_IDX;
14use crate::spec::queue::AVAIL_OFFSET_RING;
15use crate::spec::queue::DescriptorFlags;
16use crate::spec::queue::SplitDescriptor;
17use crate::spec::queue::USED_ELEMENT_SIZE;
18use crate::spec::queue::USED_OFFSET_FLAGS;
19use crate::spec::queue::USED_OFFSET_IDX;
20use crate::spec::queue::USED_OFFSET_RING;
21use crate::spec::queue::UsedElement;
22use core::mem::offset_of;
23use guestmem::GuestMemory;
24use pal_async::wait::PolledWait;
25use pal_event::Event;
26use std::time::Duration;
27
28/// Write a split virtio descriptor at the given descriptor table base.
29pub fn write_descriptor(
30    mem: &GuestMemory,
31    desc_table_base: u64,
32    index: u16,
33    addr: u64,
34    len: u32,
35    flags: DescriptorFlags,
36    next: u16,
37) {
38    let base = desc_table_base + size_of::<SplitDescriptor>() as u64 * index as u64;
39    mem.write_at(
40        base + offset_of!(SplitDescriptor, address) as u64,
41        &addr.to_le_bytes(),
42    )
43    .unwrap();
44    mem.write_at(
45        base + offset_of!(SplitDescriptor, length) as u64,
46        &len.to_le_bytes(),
47    )
48    .unwrap();
49    mem.write_at(
50        base + offset_of!(SplitDescriptor, flags_raw) as u64,
51        &u16::from(flags).to_le_bytes(),
52    )
53    .unwrap();
54    mem.write_at(
55        base + offset_of!(SplitDescriptor, next) as u64,
56        &next.to_le_bytes(),
57    )
58    .unwrap();
59}
60
61/// Initialize an avail ring (flags=0, idx=0).
62pub fn init_avail_ring(mem: &GuestMemory, avail_addr: u64) {
63    mem.write_at(avail_addr + AVAIL_OFFSET_FLAGS, &0u16.to_le_bytes())
64        .unwrap();
65    mem.write_at(avail_addr + AVAIL_OFFSET_IDX, &0u16.to_le_bytes())
66        .unwrap();
67}
68
69/// Initialize a used ring (flags=0, idx=0).
70pub fn init_used_ring(mem: &GuestMemory, used_addr: u64) {
71    mem.write_at(used_addr + USED_OFFSET_FLAGS, &0u16.to_le_bytes())
72        .unwrap();
73    mem.write_at(used_addr + USED_OFFSET_IDX, &0u16.to_le_bytes())
74        .unwrap();
75}
76
77/// Make a descriptor index available in the avail ring and bump the index.
78pub fn make_available(
79    mem: &GuestMemory,
80    avail_addr: u64,
81    queue_size: u16,
82    desc_index: u16,
83    avail_idx: &mut u16,
84) {
85    let ring_offset =
86        avail_addr + AVAIL_OFFSET_RING + AVAIL_ELEMENT_SIZE * (*avail_idx % queue_size) as u64;
87    mem.write_at(ring_offset, &desc_index.to_le_bytes())
88        .unwrap();
89    *avail_idx = avail_idx.wrapping_add(1);
90    mem.write_at(avail_addr + AVAIL_OFFSET_IDX, &avail_idx.to_le_bytes())
91        .unwrap();
92}
93
94/// Read the used ring index.
95pub fn read_used_idx(mem: &GuestMemory, used_addr: u64) -> u16 {
96    let mut buf = [0u8; 2];
97    mem.read_at(used_addr + USED_OFFSET_IDX, &mut buf).unwrap();
98    u16::from_le_bytes(buf)
99}
100
101/// Read a used ring entry (id, len) at the given ring index.
102pub fn read_used_entry(
103    mem: &GuestMemory,
104    used_addr: u64,
105    queue_size: u16,
106    index: u16,
107) -> (u32, u32) {
108    let entry_offset =
109        used_addr + USED_OFFSET_RING + USED_ELEMENT_SIZE * (index % queue_size) as u64;
110    let mut id_buf = [0u8; 4];
111    let mut len_buf = [0u8; 4];
112    mem.read_at(
113        entry_offset + offset_of!(UsedElement, id) as u64,
114        &mut id_buf,
115    )
116    .unwrap();
117    mem.read_at(
118        entry_offset + offset_of!(UsedElement, len) as u64,
119        &mut len_buf,
120    )
121    .unwrap();
122    (u32::from_le_bytes(id_buf), u32::from_le_bytes(len_buf))
123}
124
125/// Read the next used ring entry, returning `(desc_id, bytes_written)` or
126/// `None` if no new entries are available.
127///
128/// Advances `*used_idx` when an entry is consumed.
129pub fn read_used(
130    mem: &GuestMemory,
131    used_addr: u64,
132    queue_size: u16,
133    used_idx: &mut u16,
134) -> Option<(u16, u32)> {
135    let current_used_idx = read_used_idx(mem, used_addr);
136    if current_used_idx == *used_idx {
137        return None;
138    }
139    let (id, len) = read_used_entry(mem, used_addr, queue_size, *used_idx);
140    *used_idx = used_idx.wrapping_add(1);
141    Some((id as u16, len))
142}
143
144/// Wait for the next used ring entry, polling with a timeout.
145///
146/// `interrupt_event` is the event signaled by the device when it writes to
147/// the used ring. This function polls until [`read_used`] returns `Some`,
148/// panicking if 5 seconds elapse without an entry.
149pub async fn wait_for_used(
150    driver: &pal_async::DefaultDriver,
151    interrupt_event: &Event,
152    mem: &GuestMemory,
153    used_addr: u64,
154    queue_size: u16,
155    used_idx: &mut u16,
156) -> (u16, u32) {
157    let mut wait = PolledWait::new(driver, interrupt_event.clone()).unwrap();
158    mesh::CancelContext::new()
159        .with_timeout(Duration::from_secs(5))
160        .until_cancelled(async {
161            loop {
162                if let Some(entry) = read_used(mem, used_addr, queue_size, used_idx) {
163                    return entry;
164                }
165                wait.wait().await.unwrap();
166            }
167        })
168        .await
169        .expect("timed out waiting for used ring entry")
170}