Skip to main content

strat9_kernel/hardware/virtio/
rng.rs

1// VirtIO Random Number Generator Driver
2// Reference: VirtIO spec v1.2, Section 5.6 (Entropy Device)
3
4use crate::{
5    arch::x86_64::pci::{self, Bar, ProbeCriteria},
6    memory::{allocate_dma_frame, phys_to_virt},
7};
8use alloc::vec::Vec;
9use core::sync::atomic::{AtomicBool, Ordering};
10use spin::Mutex;
11
12const VIRTIO_RING_SIZE: usize = 4;
13
14pub struct VirtioRng {
15    device: VirtioDevice,
16    queue: Mutex<Virtqueue>,
17}
18
19struct VirtioDevice {
20    mmio: usize,
21}
22
23struct Virtqueue {
24    desc: *mut VirtqDesc,
25    avail: *mut VirtqAvail,
26    used: *mut VirtqUsed,
27    entropy_virt: *mut u8,
28    entropy_phys: u64,
29    free: Vec<u16>,
30    last_used_idx: u16,
31}
32
33unsafe impl Send for Virtqueue {}
34
35#[repr(C)]
36#[derive(Clone, Copy)]
37struct VirtqDesc {
38    addr: u64,
39    len: u32,
40    flags: u16,
41    next: u16,
42}
43
44#[repr(C)]
45struct VirtqAvail {
46    flags: u16,
47    idx: u16,
48    ring: [u16; VIRTIO_RING_SIZE],
49}
50
51#[repr(C)]
52struct VirtqUsed {
53    flags: u16,
54    idx: u16,
55    ring: [VirtqUsedElem; VIRTIO_RING_SIZE],
56}
57
58#[repr(C)]
59#[derive(Clone, Copy)]
60struct VirtqUsedElem {
61    id: u32,
62    len: u32,
63}
64
65const VIRTIO_F_VERSION_1: u64 = 1 << 32;
66#[allow(dead_code)]
67const VIRTIO_STATUS_RESET: u8 = 0;
68const VIRTIO_STATUS_ACKNOWLEDGE: u8 = 1;
69const VIRTIO_STATUS_DRIVER: u8 = 2;
70const VIRTIO_STATUS_DRIVER_OK: u8 = 4;
71const VIRTIO_STATUS_FEATURES_OK: u8 = 8;
72
73impl VirtioRng {
74    /// Creates a new instance.
75    pub unsafe fn new(pci_dev: pci::PciDevice) -> Result<Self, &'static str> {
76        let bar = match pci_dev.read_bar(0) {
77            Some(Bar::Memory64 { addr, .. }) => addr,
78            _ => return Err("Invalid BAR"),
79        };
80
81        let mmio = phys_to_virt(bar) as usize;
82        let mut device = VirtioDevice { mmio };
83
84        device.reset();
85        device.add_status(VIRTIO_STATUS_ACKNOWLEDGE);
86        device.add_status(VIRTIO_STATUS_DRIVER);
87
88        let features = device.read_features();
89        device.write_features(features & VIRTIO_F_VERSION_1);
90        device.add_status(VIRTIO_STATUS_FEATURES_OK);
91
92        if (device.read_status() & VIRTIO_STATUS_FEATURES_OK) == 0 {
93            return Err("Features negotiation failed");
94        }
95
96        let queue = Virtqueue::new(&mut device, 0)?;
97        device.add_status(VIRTIO_STATUS_DRIVER_OK);
98
99        Ok(Self {
100            device,
101            queue: Mutex::new(queue),
102        })
103    }
104
105    /// Reads entropy.
106    pub fn read_entropy(&self, buf: &mut [u8]) -> Result<usize, &'static str> {
107        let mut queue = self.queue.lock();
108
109        if queue.free.is_empty() {
110            return Err("No free descriptors");
111        }
112        let desc_idx = queue.free.pop().unwrap();
113
114        unsafe {
115            let desc = &mut *queue.desc.add(desc_idx as usize);
116            desc.addr = queue.entropy_phys;
117            desc.len = buf.len() as u32;
118            desc.flags = 1;
119            desc.next = 0;
120
121            let avail = &mut *queue.avail;
122            let idx = avail.idx as usize % VIRTIO_RING_SIZE;
123            avail.ring[idx] = desc_idx;
124            avail.idx = avail.idx.wrapping_add(1);
125        }
126
127        self.device.notify_queue(0);
128
129        loop {
130            unsafe {
131                let used = &*queue.used;
132                if queue.last_used_idx != used.idx {
133                    let idx = queue.last_used_idx as usize % VIRTIO_RING_SIZE;
134                    let elem = used.ring[idx];
135
136                    if elem.len as usize <= buf.len() {
137                        core::ptr::copy_nonoverlapping(
138                            queue.entropy_virt,
139                            buf.as_mut_ptr(),
140                            elem.len as usize,
141                        );
142                        queue.free.push(desc_idx);
143                        queue.last_used_idx = queue.last_used_idx.wrapping_add(1);
144                        return Ok(elem.len as usize);
145                    }
146
147                    queue.free.push(desc_idx);
148                    queue.last_used_idx = queue.last_used_idx.wrapping_add(1);
149                    return Err("Invalid entropy length");
150                }
151            }
152            core::hint::spin_loop();
153        }
154    }
155}
156
157impl VirtioDevice {
158    /// Performs the reset operation.
159    fn reset(&mut self) {
160        unsafe {
161            (self.mmio as *mut u32).write_volatile(0);
162        }
163        core::hint::spin_loop();
164    }
165
166    /// Performs the add status operation.
167    fn add_status(&mut self, status: u8) {
168        unsafe {
169            let current = ((self.mmio as *const u8).add(0x14)).read_volatile();
170            ((self.mmio as *mut u8).add(0x14)).write_volatile(current | status);
171        }
172    }
173
174    /// Reads status.
175    fn read_status(&self) -> u8 {
176        unsafe { ((self.mmio as *const u8).add(0x14)).read_volatile() }
177    }
178
179    /// Reads features.
180    fn read_features(&self) -> u64 {
181        unsafe {
182            let lo = (self.mmio as *const u32).read_volatile() as u64;
183            let hi = ((self.mmio as *const u32).add(1)).read_volatile() as u64;
184            (hi << 32) | lo
185        }
186    }
187
188    /// Writes features.
189    fn write_features(&mut self, features: u64) {
190        unsafe {
191            (self.mmio as *mut u32).write_volatile((features & 0xFFFFFFFF) as u32);
192            ((self.mmio as *mut u32).add(1)).write_volatile(((features >> 32) & 0xFFFFFFFF) as u32);
193        }
194    }
195
196    /// Performs the notify queue operation.
197    fn notify_queue(&self, queue: u16) {
198        unsafe {
199            let offset = ((self.mmio + 0x20) as *const u16).read_volatile() as usize;
200            let queue_notify = (self.mmio + 0x50 + offset * 4) as *mut u32;
201            queue_notify.write_volatile(queue as u32);
202        }
203    }
204}
205
206impl Virtqueue {
207    /// Creates a new instance.
208    fn new(device: &mut VirtioDevice, queue_idx: u16) -> Result<Self, &'static str> {
209        unsafe {
210            ((device.mmio + 0x16) as *mut u16).write_volatile(queue_idx);
211            let max_size = ((device.mmio + 0x18) as *const u16).read_volatile();
212            if max_size < VIRTIO_RING_SIZE as u16 {
213                return Err("Queue size too small");
214            }
215            ((device.mmio + 0x16) as *mut u16).write_volatile(VIRTIO_RING_SIZE as u16);
216
217            let desc_frame = allocate_dma_frame().ok_or("Failed to allocate desc")?;
218            let avail_frame = allocate_dma_frame().ok_or("Failed to allocate avail")?;
219            let used_frame = allocate_dma_frame().ok_or("Failed to allocate used")?;
220
221            let desc_phys = desc_frame.start_address.as_u64();
222            let avail_phys = avail_frame.start_address.as_u64();
223            let used_phys = used_frame.start_address.as_u64();
224
225            let desc_virt = phys_to_virt(desc_phys) as *mut VirtqDesc;
226            let avail_virt = phys_to_virt(avail_phys) as *mut VirtqAvail;
227            let used_virt = phys_to_virt(used_phys) as *mut VirtqUsed;
228
229            core::ptr::write_bytes(
230                desc_virt,
231                0,
232                VIRTIO_RING_SIZE * core::mem::size_of::<VirtqDesc>(),
233            );
234            core::ptr::write_bytes(avail_virt, 0, core::mem::size_of::<VirtqAvail>());
235            core::ptr::write_bytes(used_virt, 0, core::mem::size_of::<VirtqUsed>());
236
237            ((device.mmio + 0x10) as *mut u32).write_volatile((desc_phys & 0xFFFFFFFF) as u32);
238            ((device.mmio + 0x1A) as *mut u16).write_volatile(0xFFFF);
239
240            let buffer_frame = allocate_dma_frame().ok_or("Failed to allocate entropy buffer")?;
241            let entropy_phys = buffer_frame.start_address.as_u64();
242            let entropy_virt = phys_to_virt(entropy_phys) as *mut u8;
243            core::ptr::write_bytes(entropy_virt, 0, 4096);
244
245            let mut free = Vec::with_capacity(VIRTIO_RING_SIZE);
246            for i in 0..VIRTIO_RING_SIZE {
247                free.push(i as u16);
248            }
249
250            Ok(Self {
251                desc: desc_virt,
252                avail: avail_virt,
253                used: used_virt,
254                entropy_virt,
255                entropy_phys,
256                free,
257                last_used_idx: 0,
258            })
259        }
260    }
261}
262
263static RNG_INSTANCE: Mutex<Option<VirtioRng>> = Mutex::new(None);
264static RNG_INITIALIZED: AtomicBool = AtomicBool::new(false);
265
266/// Performs the init operation.
267pub fn init() {
268    log::info!("[VirtIO-RNG] Scanning for VirtIO RNG devices...");
269
270    let candidates = pci::probe_all(ProbeCriteria {
271        vendor_id: Some(pci::vendor::VIRTIO),
272        device_id: Some(pci::device::VIRTIO_RNG),
273        class_code: None,
274        subclass: None,
275        prog_if: None,
276    });
277
278    for pci_dev in candidates.into_iter() {
279        log::info!(
280            "VirtIO-RNG: Found device at {:?} (VEN:{:04x} DEV:{:04x})",
281            pci_dev.address,
282            pci_dev.vendor_id,
283            pci_dev.device_id
284        );
285
286        pci_dev.enable_bus_master();
287
288        match unsafe { VirtioRng::new(pci_dev) } {
289            Ok(rng) => {
290                *RNG_INSTANCE.lock() = Some(rng);
291                RNG_INITIALIZED.store(true, Ordering::SeqCst);
292                log::info!("[VirtIO-RNG] Initialized");
293                return;
294            }
295            Err(e) => {
296                log::warn!("VirtIO-RNG: Failed to initialize device: {}", e);
297            }
298        }
299    }
300
301    log::info!("[VirtIO-RNG] No device found");
302}
303
304/// Reads entropy.
305pub fn read_entropy(buf: &mut [u8]) -> Result<usize, &'static str> {
306    let rng = RNG_INSTANCE.lock();
307    match rng.as_ref() {
308        Some(rng) => rng.read_entropy(buf),
309        None => Err("RNG not initialized"),
310    }
311}
312
313/// Returns whether available.
314pub fn is_available() -> bool {
315    RNG_INITIALIZED.load(Ordering::Relaxed)
316}