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: Option<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: Some(
229                    self.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
243        fn restore(&mut self, state: Self::SavedState) -> Result<(), RestoreError> {
244            if let state::SavedState { nvram: Some(nvram) } = state {
245                self.nvram = nvram
246                    .into_iter()
247                    .map::<Result<(VariableKey, Variable), RestoreError>, _>(|e| {
248                        Ok((
249                            VariableKey {
250                                vendor: e.vendor,
251                                name: Ucs2LeVec::from_vec_with_nul(e.name)
252                                    .map_err(|e| RestoreError::InvalidSavedState(e.into()))?,
253                            },
254                            Variable {
255                                data: e.data,
256                                timestamp: e.timestamp,
257                                attr: e.attr,
258                            },
259                        ))
260                    })
261                    .collect::<Result<BTreeMap<_, _>, _>>()?;
262            }
263            Ok(())
264        }
265    }
266}
267
268/// A collection of test-implementation helpers that operate on a generic
269/// implementation of [`NvramStorage`]
270pub mod impl_agnostic_tests {
271    use crate::EFI_TIME;
272    use crate::NextVariable;
273    use crate::NvramStorage;
274    use guid::Guid;
275    use ucs2::Ucs2LeSlice;
276    use wchar::wchz;
277    use zerocopy::FromZeros;
278    use zerocopy::IntoBytes;
279
280    pub async fn test_single_variable(nvram: &mut dyn NvramStorage) {
281        let vendor = Guid::new_random();
282        let name = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var1").as_bytes()).unwrap();
283        let attr = 0x1234;
284        let data = vec![0x1, 0x2, 0x3, 0x4, 0x5];
285        let data1 = vec![0xa, 0xb, 0xc];
286        let timestamp = EFI_TIME::new_zeroed();
287
288        nvram
289            .set_variable(name, vendor, attr, data.clone(), timestamp)
290            .await
291            .unwrap();
292
293        let (result_attr, result_data, result_timestamp) =
294            nvram.get_variable(name, vendor).await.unwrap().unwrap();
295        assert_eq!(result_attr, attr);
296        assert_eq!(result_data, data);
297        assert_eq!(result_timestamp, timestamp);
298
299        let result = nvram.next_variable(Some((name, vendor))).await.unwrap();
300        assert!(matches!(result, NextVariable::EndOfList));
301
302        // set existing variable with new data
303        nvram
304            .set_variable(name, vendor, attr, data1.clone(), timestamp)
305            .await
306            .unwrap();
307
308        let (result_attr, result_data, result_timestamp) =
309            nvram.get_variable(name, vendor).await.unwrap().unwrap();
310        assert_eq!(result_attr, attr);
311        assert_eq!(result_data, data1);
312        assert_eq!(result_timestamp, timestamp);
313
314        nvram.remove_variable(name, vendor).await.unwrap();
315
316        // try to get removed variable
317        let result = nvram.get_variable(name, vendor).await.unwrap();
318        assert!(result.is_none());
319    }
320
321    pub async fn test_next(nvram: &mut dyn NvramStorage) {
322        let vendor1 = Guid::new_random();
323        let name1 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var1").as_bytes()).unwrap();
324        let vendor2 = Guid::new_random();
325        let name2 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var2").as_bytes()).unwrap();
326        let vendor3 = Guid::new_random();
327        let name3 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var3").as_bytes()).unwrap();
328        let attr = 0x1234;
329        let data = vec![0x1, 0x2, 0x3, 0x4, 0x5];
330        let timestamp = EFI_TIME::new_zeroed();
331
332        nvram
333            .set_variable(name1, vendor1, attr, data.clone(), timestamp)
334            .await
335            .unwrap();
336        nvram
337            .set_variable(name2, vendor2, attr, data.clone(), timestamp)
338            .await
339            .unwrap();
340        nvram
341            .set_variable(name3, vendor3, attr, data, timestamp)
342            .await
343            .unwrap();
344
345        let mut expected = {
346            let mut s = std::collections::BTreeSet::new();
347
348            s.insert(NextVariable::Exists {
349                name: name1.to_owned(),
350                vendor: vendor1,
351                attr,
352            });
353            s.insert(NextVariable::Exists {
354                name: name2.to_owned(),
355                vendor: vendor2,
356                attr,
357            });
358            s.insert(NextVariable::Exists {
359                name: name3.to_owned(),
360                vendor: vendor3,
361                attr,
362            });
363
364            s
365        };
366
367        let mut owned_key;
368        let mut key = None;
369        loop {
370            let var = nvram.next_variable(key).await.unwrap();
371            match &var {
372                NextVariable::InvalidKey => panic!(),
373                NextVariable::EndOfList => break,
374                NextVariable::Exists {
375                    name,
376                    vendor,
377                    attr: _,
378                } => owned_key = Some((name.clone(), *vendor)),
379            };
380
381            key = owned_key
382                .as_ref()
383                .map(|(name, vendor)| (name.as_ref(), *vendor));
384
385            let removed = expected.remove(&var);
386            assert!(removed);
387        }
388
389        assert!(expected.is_empty());
390
391        // check to make sure calls to next_variable are idempotent
392
393        let var1 = nvram.next_variable(None).await.unwrap();
394        let var2 = nvram.next_variable(None).await.unwrap();
395        assert_eq!(var1, var2);
396
397        let key = match nvram.next_variable(None).await.unwrap() {
398            NextVariable::Exists {
399                name,
400                vendor,
401                attr: _,
402            } => Some((name, vendor)),
403            _ => panic!(),
404        };
405        let key = key.as_ref().map(|(name, vendor)| (name.as_ref(), *vendor));
406
407        let var1 = nvram.next_variable(key).await.unwrap();
408        let var2 = nvram.next_variable(key).await.unwrap();
409        assert_eq!(var1, var2);
410    }
411
412    pub async fn test_multiple_variable(nvram: &mut dyn NvramStorage) {
413        let vendor1 = Guid::new_random();
414        let name1 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var1").as_bytes()).unwrap();
415        let vendor2 = Guid::new_random();
416        let name2 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var2").as_bytes()).unwrap();
417        let vendor3 = Guid::new_random();
418        let name3 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var3").as_bytes()).unwrap();
419        let attr = 0x1234;
420        let data = vec![0x1, 0x2, 0x3, 0x4, 0x5];
421        let timestamp = EFI_TIME::new_zeroed();
422
423        // add all variables to nvram
424        nvram
425            .set_variable(name1, vendor1, attr, data.clone(), timestamp)
426            .await
427            .unwrap();
428        nvram
429            .set_variable(name2, vendor2, attr, data.clone(), timestamp)
430            .await
431            .unwrap();
432        nvram
433            .set_variable(name3, vendor3, attr, data.clone(), timestamp)
434            .await
435            .unwrap();
436
437        let (result_attr, result_data, result_timestamp) =
438            nvram.get_variable(name1, vendor1).await.unwrap().unwrap();
439        assert_eq!(result_attr, attr);
440        assert_eq!(result_data, data);
441        assert_eq!(result_timestamp, timestamp);
442
443        let (result_attr, result_data, result_timestamp) =
444            nvram.get_variable(name2, vendor2).await.unwrap().unwrap();
445        assert_eq!(result_attr, attr);
446        assert_eq!(result_data, data);
447        assert_eq!(result_timestamp, timestamp);
448
449        let (result_attr, result_data, result_timestamp) =
450            nvram.get_variable(name3, vendor3).await.unwrap().unwrap();
451        assert_eq!(result_attr, attr);
452        assert_eq!(result_data, data);
453        assert_eq!(result_timestamp, timestamp);
454
455        // throw an append in there for good measure
456        let appended = nvram
457            .append_variable(name1, vendor1, vec![6, 7, 8], timestamp)
458            .await
459            .unwrap();
460        assert!(appended);
461
462        let (result_attr, result_data, result_timestamp) =
463            nvram.get_variable(name1, vendor1).await.unwrap().unwrap();
464        assert_eq!(result_attr, attr);
465        assert_eq!(result_data, (1..=8).collect::<Vec<u8>>());
466        assert_eq!(result_timestamp, timestamp);
467    }
468}
469
470#[cfg(test)]
471mod tests {
472    use super::impl_agnostic_tests;
473    use super::*;
474    use pal_async::async_test;
475
476    #[async_test]
477    async fn nvram_trait_single_variable() {
478        let mut nvram = InMemoryNvram::new();
479        impl_agnostic_tests::test_single_variable(&mut nvram).await;
480    }
481
482    #[async_test]
483    async fn nvram_trait_next() {
484        let mut nvram = InMemoryNvram::new();
485        impl_agnostic_tests::test_next(&mut nvram).await;
486    }
487
488    #[async_test]
489    async fn nvram_trait_multiple_variable() {
490        let mut nvram = InMemoryNvram::new();
491        impl_agnostic_tests::test_multiple_variable(&mut nvram).await;
492    }
493}
uefi_nvram_storage/in_memory.rs

uefi_nvram_storage/
in_memory.rs
