1use 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
28const 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
94pub struct Rng {
96 worker_thread: Option<WorkerThread<Worker>>,
97 virtio_features: u64,
98}
99
100impl Rng {
101 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 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}