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

strat9_kernel/hardware/virtio/
rng.rs
