uefi_nvram_storage/
in_memory.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Provides an in-memory implementation of [`NvramStorage`] that doesn't
5//! automatically persist to disk.
6
7use crate::EFI_TIME;
8use crate::NextVariable;
9use crate::NvramStorage;
10use crate::NvramStorageError;
11use guid::Guid;
12use std::collections::BTreeMap;
13use std::fmt::Display;
14use ucs2::Ucs2LeSlice;
15use ucs2::Ucs2LeVec;
16
17#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
18struct VariableKey {
19    vendor: Guid,
20    name: Ucs2LeVec,
21}
22
23impl Display for VariableKey {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        write!(f, "{}-{}", self.vendor, self.name)
26    }
27}
28
29#[derive(Clone, Debug)]
30#[cfg_attr(feature = "inspect", derive(inspect::Inspect))]
31struct Variable {
32    data: Vec<u8>,
33    timestamp: EFI_TIME,
34    #[cfg_attr(feature = "inspect", inspect(hex))]
35    attr: u32,
36}
37
38/// An in-memory implementation of [`NvramStorage`].
39#[derive(Debug)]
40#[cfg_attr(feature = "inspect", derive(inspect::Inspect))]
41pub struct InMemoryNvram {
42    #[cfg_attr(feature = "inspect", inspect(iter_by_key))]
43    nvram: BTreeMap<VariableKey, Variable>,
44}
45
46pub struct VariableEntry<'a> {
47    pub vendor: Guid,
48    pub name: &'a Ucs2LeSlice,
49    pub data: &'a [u8],
50    pub timestamp: EFI_TIME,
51    pub attr: u32,
52}
53
54impl InMemoryNvram {
55    pub fn new() -> Self {
56        Self {
57            nvram: Default::default(),
58        }
59    }
60
61    pub fn iter(&self) -> impl Iterator<Item = VariableEntry<'_>> {
62        self.nvram.iter().map(|(k, v)| VariableEntry {
63            vendor: k.vendor,
64            name: k.name.as_ref(),
65            data: v.data.as_slice(),
66            timestamp: v.timestamp,
67            attr: v.attr,
68        })
69    }
70
71    pub fn clear(&mut self) {
72        self.nvram.clear()
73    }
74}
75
76#[async_trait::async_trait]
77impl NvramStorage for InMemoryNvram {
78    async fn get_variable(
79        &mut self,
80        name: &Ucs2LeSlice,
81        vendor: Guid,
82    ) -> Result<Option<(u32, Vec<u8>, EFI_TIME)>, NvramStorageError> {
83        Ok(self
84            .nvram
85            .get(&VariableKey {
86                vendor,
87                name: name.to_ucs2_le_vec(),
88            })
89            .map(|v| (v.attr, v.data.clone(), v.timestamp)))
90    }
91
92    async fn set_variable(
93        &mut self,
94        name: &Ucs2LeSlice,
95        vendor: Guid,
96        attr: u32,
97        data: Vec<u8>,
98        timestamp: EFI_TIME,
99    ) -> Result<(), NvramStorageError> {
100        self.nvram.insert(
101            VariableKey {
102                vendor,
103                name: name.to_ucs2_le_vec(),
104            },
105            Variable {
106                data,
107                timestamp,
108                attr,
109            },
110        );
111        Ok(())
112    }
113
114    async fn append_variable(
115        &mut self,
116        name: &Ucs2LeSlice,
117        vendor: Guid,
118        data: Vec<u8>,
119        timestamp: EFI_TIME,
120    ) -> Result<bool, NvramStorageError> {
121        match self.nvram.get_mut(&VariableKey {
122            vendor,
123            name: name.to_ucs2_le_vec(),
124        }) {
125            Some(val) => {
126                val.data.extend_from_slice(&data);
127                val.timestamp = timestamp;
128                Ok(true)
129            }
130            None => Ok(false),
131        }
132    }
133
134    async fn remove_variable(
135        &mut self,
136        name: &Ucs2LeSlice,
137        vendor: Guid,
138    ) -> Result<bool, NvramStorageError> {
139        Ok(self
140            .nvram
141            .remove(&VariableKey {
142                vendor,
143                name: name.to_ucs2_le_vec(),
144            })
145            .is_some())
146    }
147
148    async fn next_variable(
149        &mut self,
150        name_vendor: Option<(&Ucs2LeSlice, Guid)>,
151    ) -> Result<NextVariable, NvramStorageError> {
152        let key = &name_vendor.map(|(name, vendor)| VariableKey {
153            vendor,
154            name: name.to_ucs2_le_vec(),
155        });
156
157        if let Some(key) = key {
158            let mut range = self.nvram.range(key..);
159            if let Some((found_key, _)) = range.next() {
160                if found_key == key {
161                    Ok(match range.next() {
162                        Some(v) => NextVariable::Exists {
163                            name: v.0.name.clone(),
164                            vendor: v.0.vendor,
165                            attr: v.1.attr,
166                        },
167                        None => NextVariable::EndOfList,
168                    })
169                } else {
170                    Ok(NextVariable::InvalidKey)
171                }
172            } else {
173                Ok(NextVariable::EndOfList)
174            }
175        } else {
176            Ok(match self.nvram.iter().next() {
177                Some(v) => NextVariable::Exists {
178                    name: v.0.name.clone(),
179                    vendor: v.0.vendor,
180                    attr: v.1.attr,
181                },
182                None => NextVariable::EndOfList,
183            })
184        }
185    }
186}
187
188#[cfg(feature = "save_restore")]
189mod save_restore {
190    use super::*;
191    use vmcore::save_restore::RestoreError;
192    use vmcore::save_restore::SaveError;
193    use vmcore::save_restore::SaveRestore;
194
195    mod state {
196        use guid::Guid;
197        use uefi_specs::uefi::time::EFI_TIME;
198        use vmcore::save_restore::SavedStateRoot;
199
200        #[derive(mesh_protobuf::Protobuf)]
201        #[mesh(package = "nvram_entry")]
202        pub struct NvramEntry {
203            #[mesh(1)]
204            pub vendor: Guid,
205            #[mesh(2)]
206            pub name: Vec<u8>,
207            #[mesh(3)]
208            pub data: Vec<u8>,
209            #[mesh(4, encoding = "mesh_protobuf::encoding::ZeroCopyEncoding")]
210            pub timestamp: EFI_TIME,
211            #[mesh(5)]
212            pub attr: u32,
213        }
214
215        #[derive(mesh_protobuf::Protobuf, SavedStateRoot)]
216        #[mesh(package = "firmware.in_memory_nvram")]
217        pub struct SavedState {
218            #[mesh(1)]
219            pub nvram: Vec<NvramEntry>,
220        }
221    }
222
223    impl SaveRestore for InMemoryNvram {
224        type SavedState = state::SavedState;
225
226        fn save(&mut self) -> Result<Self::SavedState, SaveError> {
227            Ok(state::SavedState {
228                nvram: self
229                    .nvram
230                    .iter()
231                    .map(|(k, v)| state::NvramEntry {
232                        vendor: k.vendor,
233                        name: k.name.clone().into_inner(),
234                        data: v.data.clone(),
235                        timestamp: v.timestamp,
236                        attr: v.attr,
237                    })
238                    .collect(),
239            })
240        }
241
242        fn restore(&mut self, state: Self::SavedState) -> Result<(), RestoreError> {
243            let state::SavedState { nvram } = state;
244            self.nvram = nvram
245                .into_iter()
246                .map::<Result<(VariableKey, Variable), RestoreError>, _>(|e| {
247                    Ok((
248                        VariableKey {
249                            vendor: e.vendor,
250                            name: Ucs2LeVec::from_vec_with_nul(e.name)
251                                .map_err(|e| RestoreError::InvalidSavedState(e.into()))?,
252                        },
253                        Variable {
254                            data: e.data,
255                            timestamp: e.timestamp,
256                            attr: e.attr,
257                        },
258                    ))
259                })
260                .collect::<Result<BTreeMap<_, _>, _>>()?;
261            Ok(())
262        }
263    }
264}
265
266/// A collection of test-implementation helpers that operate on a generic
267/// implementation of [`NvramStorage`]
268pub mod impl_agnostic_tests {
269    use crate::EFI_TIME;
270    use crate::NextVariable;
271    use crate::NvramStorage;
272    use guid::Guid;
273    use ucs2::Ucs2LeSlice;
274    use wchar::wchz;
275    use zerocopy::FromZeros;
276    use zerocopy::IntoBytes;
277
278    pub async fn test_single_variable(nvram: &mut dyn NvramStorage) {
279        let vendor = Guid::new_random();
280        let name = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var1").as_bytes()).unwrap();
281        let attr = 0x1234;
282        let data = vec![0x1, 0x2, 0x3, 0x4, 0x5];
283        let data1 = vec![0xa, 0xb, 0xc];
284        let timestamp = EFI_TIME::new_zeroed();
285
286        nvram
287            .set_variable(name, vendor, attr, data.clone(), timestamp)
288            .await
289            .unwrap();
290
291        let (result_attr, result_data, result_timestamp) =
292            nvram.get_variable(name, vendor).await.unwrap().unwrap();
293        assert_eq!(result_attr, attr);
294        assert_eq!(result_data, data);
295        assert_eq!(result_timestamp, timestamp);
296
297        let result = nvram.next_variable(Some((name, vendor))).await.unwrap();
298        assert!(matches!(result, NextVariable::EndOfList));
299
300        // set existing variable with new data
301        nvram
302            .set_variable(name, vendor, attr, data1.clone(), timestamp)
303            .await
304            .unwrap();
305
306        let (result_attr, result_data, result_timestamp) =
307            nvram.get_variable(name, vendor).await.unwrap().unwrap();
308        assert_eq!(result_attr, attr);
309        assert_eq!(result_data, data1);
310        assert_eq!(result_timestamp, timestamp);
311
312        nvram.remove_variable(name, vendor).await.unwrap();
313
314        // try to get removed variable
315        let result = nvram.get_variable(name, vendor).await.unwrap();
316        assert!(result.is_none());
317    }
318
319    pub async fn test_next(nvram: &mut dyn NvramStorage) {
320        let vendor1 = Guid::new_random();
321        let name1 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var1").as_bytes()).unwrap();
322        let vendor2 = Guid::new_random();
323        let name2 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var2").as_bytes()).unwrap();
324        let vendor3 = Guid::new_random();
325        let name3 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var3").as_bytes()).unwrap();
326        let attr = 0x1234;
327        let data = vec![0x1, 0x2, 0x3, 0x4, 0x5];
328        let timestamp = EFI_TIME::new_zeroed();
329
330        nvram
331            .set_variable(name1, vendor1, attr, data.clone(), timestamp)
332            .await
333            .unwrap();
334        nvram
335            .set_variable(name2, vendor2, attr, data.clone(), timestamp)
336            .await
337            .unwrap();
338        nvram
339            .set_variable(name3, vendor3, attr, data, timestamp)
340            .await
341            .unwrap();
342
343        let mut expected = {
344            let mut s = std::collections::BTreeSet::new();
345
346            s.insert(NextVariable::Exists {
347                name: name1.to_owned(),
348                vendor: vendor1,
349                attr,
350            });
351            s.insert(NextVariable::Exists {
352                name: name2.to_owned(),
353                vendor: vendor2,
354                attr,
355            });
356            s.insert(NextVariable::Exists {
357                name: name3.to_owned(),
358                vendor: vendor3,
359                attr,
360            });
361
362            s
363        };
364
365        let mut owned_key;
366        let mut key = None;
367        loop {
368            let var = nvram.next_variable(key).await.unwrap();
369            match &var {
370                NextVariable::InvalidKey => panic!(),
371                NextVariable::EndOfList => break,
372                NextVariable::Exists {
373                    name,
374                    vendor,
375                    attr: _,
376                } => owned_key = Some((name.clone(), *vendor)),
377            };
378
379            key = owned_key
380                .as_ref()
381                .map(|(name, vendor)| (name.as_ref(), *vendor));
382
383            let removed = expected.remove(&var);
384            assert!(removed);
385        }
386
387        assert!(expected.is_empty());
388
389        // check to make sure calls to next_variable are idempotent
390
391        let var1 = nvram.next_variable(None).await.unwrap();
392        let var2 = nvram.next_variable(None).await.unwrap();
393        assert_eq!(var1, var2);
394
395        let key = match nvram.next_variable(None).await.unwrap() {
396            NextVariable::Exists {
397                name,
398                vendor,
399                attr: _,
400            } => Some((name, vendor)),
401            _ => panic!(),
402        };
403        let key = key.as_ref().map(|(name, vendor)| (name.as_ref(), *vendor));
404
405        let var1 = nvram.next_variable(key).await.unwrap();
406        let var2 = nvram.next_variable(key).await.unwrap();
407        assert_eq!(var1, var2);
408    }
409
410    pub async fn test_multiple_variable(nvram: &mut dyn NvramStorage) {
411        let vendor1 = Guid::new_random();
412        let name1 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var1").as_bytes()).unwrap();
413        let vendor2 = Guid::new_random();
414        let name2 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var2").as_bytes()).unwrap();
415        let vendor3 = Guid::new_random();
416        let name3 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var3").as_bytes()).unwrap();
417        let attr = 0x1234;
418        let data = vec![0x1, 0x2, 0x3, 0x4, 0x5];
419        let timestamp = EFI_TIME::new_zeroed();
420
421        // add all variables to nvram
422        nvram
423            .set_variable(name1, vendor1, attr, data.clone(), timestamp)
424            .await
425            .unwrap();
426        nvram
427            .set_variable(name2, vendor2, attr, data.clone(), timestamp)
428            .await
429            .unwrap();
430        nvram
431            .set_variable(name3, vendor3, attr, data.clone(), timestamp)
432            .await
433            .unwrap();
434
435        let (result_attr, result_data, result_timestamp) =
436            nvram.get_variable(name1, vendor1).await.unwrap().unwrap();
437        assert_eq!(result_attr, attr);
438        assert_eq!(result_data, data);
439        assert_eq!(result_timestamp, timestamp);
440
441        let (result_attr, result_data, result_timestamp) =
442            nvram.get_variable(name2, vendor2).await.unwrap().unwrap();
443        assert_eq!(result_attr, attr);
444        assert_eq!(result_data, data);
445        assert_eq!(result_timestamp, timestamp);
446
447        let (result_attr, result_data, result_timestamp) =
448            nvram.get_variable(name3, vendor3).await.unwrap().unwrap();
449        assert_eq!(result_attr, attr);
450        assert_eq!(result_data, data);
451        assert_eq!(result_timestamp, timestamp);
452
453        // throw an append in there for good measure
454        let appended = nvram
455            .append_variable(name1, vendor1, vec![6, 7, 8], timestamp)
456            .await
457            .unwrap();
458        assert!(appended);
459
460        let (result_attr, result_data, result_timestamp) =
461            nvram.get_variable(name1, vendor1).await.unwrap().unwrap();
462        assert_eq!(result_attr, attr);
463        assert_eq!(result_data, (1..=8).collect::<Vec<u8>>());
464        assert_eq!(result_timestamp, timestamp);
465    }
466}
467
468#[cfg(test)]
469mod tests {
470    use super::impl_agnostic_tests;
471    use super::*;
472    use pal_async::async_test;
473
474    #[async_test]
475    async fn nvram_trait_single_variable() {
476        let mut nvram = InMemoryNvram::new();
477        impl_agnostic_tests::test_single_variable(&mut nvram).await;
478    }
479
480    #[async_test]
481    async fn nvram_trait_next() {
482        let mut nvram = InMemoryNvram::new();
483        impl_agnostic_tests::test_next(&mut nvram).await;
484    }
485
486    #[async_test]
487    async fn nvram_trait_multiple_variable() {
488        let mut nvram = InMemoryNvram::new();
489        impl_agnostic_tests::test_multiple_variable(&mut nvram).await;
490    }
491}
uefi_nvram_storage/in_memory.rs

uefi_nvram_storage/
in_memory.rs
