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