devices/virtio/
rng.rs

1// Copyright 2017 The ChromiumOS Authors
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use std::collections::BTreeMap;
6use std::io::Write;
7
8use anyhow::anyhow;
9use anyhow::Context;
10use base::error;
11use base::warn;
12use base::Event;
13use base::EventToken;
14use base::RawDescriptor;
15use base::WaitContext;
16use base::WorkerThread;
17use rand::rngs::OsRng;
18use rand::RngCore;
19use snapshot::AnySnapshot;
20use vm_memory::GuestMemory;
21
22use super::DeviceType;
23use super::Interrupt;
24use super::Queue;
25use super::VirtioDevice;
26
27const QUEUE_SIZE: u16 = 256;
28const QUEUE_SIZES: &[u16] = &[QUEUE_SIZE];
29
30// Chosen to match the Linux guest driver RNG buffer refill size.
31const CHUNK_SIZE: usize = 64;
32
33struct Worker {
34    queue: Queue,
35}
36
37impl Worker {
38    fn process_queue(&mut self) {
39        let mut rand_bytes = [0u8; CHUNK_SIZE];
40        let mut needs_interrupt = false;
41
42        while let Some(mut avail_desc) = self.queue.pop() {
43            let writer = &mut avail_desc.writer;
44            while writer.available_bytes() > 0 {
45                let chunk_size = writer.available_bytes().min(CHUNK_SIZE);
46                let chunk = &mut rand_bytes[..chunk_size];
47                OsRng.fill_bytes(chunk);
48                if let Err(e) = writer.write_all(chunk) {
49                    warn!("Failed to write random data to the guest: {}", e);
50                    break;
51                }
52            }
53
54            self.queue.add_used(avail_desc);
55            needs_interrupt = true;
56        }
57
58        if needs_interrupt {
59            self.queue.trigger_interrupt();
60        }
61    }
62
63    fn run(&mut self, kill_evt: Event) -> anyhow::Result<()> {
64        #[derive(EventToken)]
65        enum Token {
66            QueueAvailable,
67            Kill,
68        }
69
70        let wait_ctx = WaitContext::build_with(&[
71            (self.queue.event(), Token::QueueAvailable),
72            (&kill_evt, Token::Kill),
73        ])
74        .context("failed creating WaitContext")?;
75
76        let mut exiting = false;
77        while !exiting {
78            let events = wait_ctx.wait().context("failed polling for events")?;
79            for event in events.iter().filter(|e| e.is_readable) {
80                match event.token {
81                    Token::QueueAvailable => {
82                        self.queue
83                            .event()
84                            .wait()
85                            .context("failed reading queue Event")?;
86                        self.process_queue();
87                    }
88                    Token::Kill => exiting = true,
89                }
90            }
91        }
92
93        Ok(())
94    }
95}
96
97/// Virtio device for exposing entropy to the guest OS through virtio.
98pub struct Rng {
99    worker_thread: Option<WorkerThread<Worker>>,
100    virtio_features: u64,
101}
102
103impl Rng {
104    /// Create a new virtio rng device that gets random data from /dev/urandom.
105    pub fn new(virtio_features: u64) -> anyhow::Result<Rng> {
106        Ok(Rng {
107            worker_thread: None,
108            virtio_features,
109        })
110    }
111}
112
113impl VirtioDevice for Rng {
114    fn keep_rds(&self) -> Vec<RawDescriptor> {
115        Vec::new()
116    }
117
118    fn device_type(&self) -> DeviceType {
119        DeviceType::Rng
120    }
121
122    fn queue_max_sizes(&self) -> &[u16] {
123        QUEUE_SIZES
124    }
125
126    fn features(&self) -> u64 {
127        self.virtio_features
128    }
129
130    fn activate(
131        &mut self,
132        _mem: GuestMemory,
133        _interrupt: Interrupt,
134        mut queues: BTreeMap<usize, Queue>,
135    ) -> anyhow::Result<()> {
136        if queues.len() != 1 {
137            return Err(anyhow!("expected 1 queue, got {}", queues.len()));
138        }
139
140        let queue = queues.remove(&0).unwrap();
141
142        self.worker_thread = Some(WorkerThread::start("v_rng", move |kill_evt| {
143            let mut worker = Worker { queue };
144            if let Err(e) = worker.run(kill_evt) {
145                error!("rng worker thread failed: {:#}", e);
146            }
147            worker
148        }));
149
150        Ok(())
151    }
152
153    fn reset(&mut self) -> anyhow::Result<()> {
154        if let Some(worker_thread) = self.worker_thread.take() {
155            let _worker = worker_thread.stop();
156        }
157        Ok(())
158    }
159
160    fn virtio_sleep(&mut self) -> anyhow::Result<Option<BTreeMap<usize, Queue>>> {
161        if let Some(worker_thread) = self.worker_thread.take() {
162            let worker = worker_thread.stop();
163            return Ok(Some(BTreeMap::from([(0, worker.queue)])));
164        }
165        Ok(None)
166    }
167
168    fn virtio_wake(
169        &mut self,
170        queues_state: Option<(GuestMemory, Interrupt, BTreeMap<usize, Queue>)>,
171    ) -> anyhow::Result<()> {
172        if let Some((mem, interrupt, queues)) = queues_state {
173            self.activate(mem, interrupt, queues)?;
174        }
175        Ok(())
176    }
177
178    fn virtio_snapshot(&mut self) -> anyhow::Result<AnySnapshot> {
179        // `virtio_sleep` ensures there is no pending state, except for the `Queue`s, which are
180        // handled at a higher layer.
181        AnySnapshot::to_any(())
182    }
183
184    fn virtio_restore(&mut self, data: AnySnapshot) -> anyhow::Result<()> {
185        let () = AnySnapshot::from_any(data)?;
186        Ok(())
187    }
188}
devices/virtio/rng.rs

devices/virtio/
rng.rs
