vmm/devices/virtio/rng/
persist.rs

1// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Defines the structures needed for saving/restoring entropy devices.
5
6use std::sync::Arc;
7
8use serde::{Deserialize, Serialize};
9
10use crate::devices::virtio::generated::virtio_ids::VIRTIO_ID_RNG;
11use crate::devices::virtio::persist::{PersistError as VirtioStateError, VirtioDeviceState};
12use crate::devices::virtio::queue::FIRECRACKER_MAX_QUEUE_SIZE;
13use crate::devices::virtio::rng::{Entropy, EntropyError, RNG_NUM_QUEUES};
14use crate::devices::virtio::transport::VirtioInterrupt;
15use crate::rate_limiter::RateLimiter;
16use crate::rate_limiter::persist::RateLimiterState;
17use crate::snapshot::Persist;
18use crate::vstate::memory::GuestMemoryMmap;
19
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct EntropyState {
22    pub virtio_state: VirtioDeviceState,
23    rate_limiter_state: RateLimiterState,
24}
25
26#[derive(Debug)]
27pub struct EntropyConstructorArgs {
28    pub mem: GuestMemoryMmap,
29}
30
31#[derive(Debug, thiserror::Error, displaydoc::Display)]
32pub enum EntropyPersistError {
33    /// Create entropy: {0}
34    CreateEntropy(#[from] EntropyError),
35    /// Virtio state: {0}
36    VirtioState(#[from] VirtioStateError),
37    /// Restore rate limiter: {0}
38    RestoreRateLimiter(#[from] std::io::Error),
39}
40
41impl Persist<'_> for Entropy {
42    type State = EntropyState;
43    type ConstructorArgs = EntropyConstructorArgs;
44    type Error = EntropyPersistError;
45
46    fn save(&self) -> Self::State {
47        EntropyState {
48            virtio_state: VirtioDeviceState::from_device(self),
49            rate_limiter_state: self.rate_limiter().save(),
50        }
51    }
52
53    fn restore(
54        constructor_args: Self::ConstructorArgs,
55        state: &Self::State,
56    ) -> Result<Self, Self::Error> {
57        let queues = state.virtio_state.build_queues_checked(
58            &constructor_args.mem,
59            VIRTIO_ID_RNG,
60            RNG_NUM_QUEUES,
61            FIRECRACKER_MAX_QUEUE_SIZE,
62        )?;
63
64        let rate_limiter = RateLimiter::restore((), &state.rate_limiter_state)?;
65        let mut entropy = Entropy::new_with_queues(queues, rate_limiter)?;
66        entropy.set_avail_features(state.virtio_state.avail_features);
67        entropy.set_acked_features(state.virtio_state.acked_features);
68
69        Ok(entropy)
70    }
71}
72
73#[cfg(test)]
74mod tests {
75
76    use super::*;
77    use crate::devices::virtio::device::VirtioDevice;
78    use crate::devices::virtio::rng::device::ENTROPY_DEV_ID;
79    use crate::devices::virtio::test_utils::default_interrupt;
80    use crate::devices::virtio::test_utils::test::create_virtio_mem;
81    use crate::snapshot::Snapshot;
82
83    #[test]
84    fn test_persistence() {
85        let mut mem = vec![0u8; 4096];
86        let entropy = Entropy::new(RateLimiter::default()).unwrap();
87
88        Snapshot::new(entropy.save())
89            .save(&mut mem.as_mut_slice())
90            .unwrap();
91
92        let guest_mem = create_virtio_mem();
93        let restored = Entropy::restore(
94            EntropyConstructorArgs { mem: guest_mem },
95            &Snapshot::load_without_crc_check(mem.as_slice())
96                .unwrap()
97                .data,
98        )
99        .unwrap();
100
101        assert_eq!(restored.device_type(), VIRTIO_ID_RNG);
102        assert_eq!(restored.id(), ENTROPY_DEV_ID);
103        assert!(!restored.is_activated());
104        assert!(!entropy.is_activated());
105        assert_eq!(restored.avail_features(), entropy.avail_features());
106        assert_eq!(restored.acked_features(), entropy.acked_features());
107    }
108}
vmm/devices/virtio/rng/persist.rs

vmm/devices/virtio/rng/
persist.rs
