supermachine 0.4.13

// virtio-rng — single-queue device that the guest's hw_random core
// uses to reseed its CRNG. Linux's kernel polls it during boot and
// whenever the CRNG entropy estimate dips. Without this device the
// guest pulls entropy from the timer-jitter / ARMv8 RNDR / other
// fallbacks, which can stall boot for hundreds of milliseconds.
//
// One RX queue: the guest provides writable buffer chains; we fill
// each with bytes from /dev/urandom and push back as used.

use std::io::Read;
use std::sync::atomic::AtomicBool;
use std::sync::{Arc, Mutex};

use super::queue::{Queue, VRING_DESC_F_WRITE};
use super::{VirtioDevice, VIRTIO_ID_RNG};

pub struct VirtioRng {
    queues: Mutex<Vec<Queue>>,
    activated: AtomicBool,
    irq_raise: Mutex<Option<Arc<dyn Fn() + Send + Sync>>>,
}

impl VirtioRng {
    pub fn new() -> Self {
        Self {
            queues: Mutex::new(Vec::new()),
            activated: AtomicBool::new(false),
            irq_raise: Mutex::new(None),
        }
    }

    pub fn set_irq_raise(&self, f: Arc<dyn Fn() + Send + Sync>) {
        *self.irq_raise.lock().unwrap() = Some(f);
    }

    fn drain(&self) {
        if !self.activated.load(std::sync::atomic::Ordering::Acquire) {
            return;
        }
        let mut qs = self.queues.lock().unwrap();
        let q = match qs.get_mut(0) {
            Some(q) => q,
            None => return,
        };
        if !q.ready {
            return;
        }
        let mut any = false;
        let mut urandom = match std::fs::File::open("/dev/urandom") {
            Ok(f) => f,
            Err(_) => return,
        };
        loop {
            let (head, chain) = match q.pop_chain() {
                Some(p) => p,
                None => break,
            };
            let mut written: u32 = 0;
            let mut buf = [0u8; 4096];
            for d in &chain {
                if d.flags & VRING_DESC_F_WRITE == 0 {
                    continue;
                }
                let mut remaining = d.len as usize;
                let mut off = 0u64;
                while remaining > 0 {
                    let take = remaining.min(buf.len());
                    if urandom.read_exact(&mut buf[..take]).is_err() {
                        break;
                    }
                    q.mem.write_slice(d.addr + off, &buf[..take]);
                    written += take as u32;
                    off += take as u64;
                    remaining -= take;
                }
            }
            q.add_used(head, written);
            any = true;
        }
        drop(qs);
        if any {
            if let Some(f) = self.irq_raise.lock().unwrap().clone() {
                f();
            }
        }
    }
}

impl VirtioDevice for VirtioRng {
    fn device_id(&self) -> u32 {
        VIRTIO_ID_RNG
    }
    fn num_queues(&self) -> usize {
        1
    }
    fn features(&self) -> u64 {
        1u64 << 32 /* VIRTIO_F_VERSION_1 */
    }
    fn notify(&self, _q: u16) {
        self.drain();
    }
    fn activate(&self, queues: Vec<Queue>) {
        *self.queues.lock().unwrap() = queues;
        self.activated
            .store(true, std::sync::atomic::Ordering::Release);
        eprintln!("[virtio-rng] activated");
    }
    fn snapshot_queues(&self) -> Vec<Queue> {
        self.queues.lock().unwrap().clone()
    }
}