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/// A collection of test-implementation helpers that operate on a generic
189/// implementation of [`NvramStorage`]
190pub mod impl_agnostic_tests {
191    use crate::EFI_TIME;
192    use crate::NextVariable;
193    use crate::NvramStorage;
194    use guid::Guid;
195    use ucs2::Ucs2LeSlice;
196    use wchar::wchz;
197    use zerocopy::FromZeros;
198    use zerocopy::IntoBytes;
199
200    pub async fn test_single_variable(nvram: &mut dyn NvramStorage) {
201        let vendor = Guid::new_random();
202        let name = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var1").as_bytes()).unwrap();
203        let attr = 0x1234;
204        let data = vec![0x1, 0x2, 0x3, 0x4, 0x5];
205        let data1 = vec![0xa, 0xb, 0xc];
206        let timestamp = EFI_TIME::new_zeroed();
207
208        nvram
209            .set_variable(name, vendor, attr, data.clone(), timestamp)
210            .await
211            .unwrap();
212
213        let (result_attr, result_data, result_timestamp) =
214            nvram.get_variable(name, vendor).await.unwrap().unwrap();
215        assert_eq!(result_attr, attr);
216        assert_eq!(result_data, data);
217        assert_eq!(result_timestamp, timestamp);
218
219        let result = nvram.next_variable(Some((name, vendor))).await.unwrap();
220        assert!(matches!(result, NextVariable::EndOfList));
221
222        // set existing variable with new data
223        nvram
224            .set_variable(name, vendor, attr, data1.clone(), timestamp)
225            .await
226            .unwrap();
227
228        let (result_attr, result_data, result_timestamp) =
229            nvram.get_variable(name, vendor).await.unwrap().unwrap();
230        assert_eq!(result_attr, attr);
231        assert_eq!(result_data, data1);
232        assert_eq!(result_timestamp, timestamp);
233
234        nvram.remove_variable(name, vendor).await.unwrap();
235
236        // try to get removed variable
237        let result = nvram.get_variable(name, vendor).await.unwrap();
238        assert!(result.is_none());
239    }
240
241    pub async fn test_next(nvram: &mut dyn NvramStorage) {
242        let vendor1 = Guid::new_random();
243        let name1 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var1").as_bytes()).unwrap();
244        let vendor2 = Guid::new_random();
245        let name2 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var2").as_bytes()).unwrap();
246        let vendor3 = Guid::new_random();
247        let name3 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var3").as_bytes()).unwrap();
248        let attr = 0x1234;
249        let data = vec![0x1, 0x2, 0x3, 0x4, 0x5];
250        let timestamp = EFI_TIME::new_zeroed();
251
252        nvram
253            .set_variable(name1, vendor1, attr, data.clone(), timestamp)
254            .await
255            .unwrap();
256        nvram
257            .set_variable(name2, vendor2, attr, data.clone(), timestamp)
258            .await
259            .unwrap();
260        nvram
261            .set_variable(name3, vendor3, attr, data, timestamp)
262            .await
263            .unwrap();
264
265        let mut expected = {
266            let mut s = std::collections::BTreeSet::new();
267
268            s.insert(NextVariable::Exists {
269                name: name1.to_owned(),
270                vendor: vendor1,
271                attr,
272            });
273            s.insert(NextVariable::Exists {
274                name: name2.to_owned(),
275                vendor: vendor2,
276                attr,
277            });
278            s.insert(NextVariable::Exists {
279                name: name3.to_owned(),
280                vendor: vendor3,
281                attr,
282            });
283
284            s
285        };
286
287        let mut owned_key;
288        let mut key = None;
289        loop {
290            let var = nvram.next_variable(key).await.unwrap();
291            match &var {
292                NextVariable::InvalidKey => panic!(),
293                NextVariable::EndOfList => break,
294                NextVariable::Exists {
295                    name,
296                    vendor,
297                    attr: _,
298                } => owned_key = Some((name.clone(), *vendor)),
299            };
300
301            key = owned_key
302                .as_ref()
303                .map(|(name, vendor)| (name.as_ref(), *vendor));
304
305            let removed = expected.remove(&var);
306            assert!(removed);
307        }
308
309        assert!(expected.is_empty());
310
311        // check to make sure calls to next_variable are idempotent
312
313        let var1 = nvram.next_variable(None).await.unwrap();
314        let var2 = nvram.next_variable(None).await.unwrap();
315        assert_eq!(var1, var2);
316
317        let key = match nvram.next_variable(None).await.unwrap() {
318            NextVariable::Exists {
319                name,
320                vendor,
321                attr: _,
322            } => Some((name, vendor)),
323            _ => panic!(),
324        };
325        let key = key.as_ref().map(|(name, vendor)| (name.as_ref(), *vendor));
326
327        let var1 = nvram.next_variable(key).await.unwrap();
328        let var2 = nvram.next_variable(key).await.unwrap();
329        assert_eq!(var1, var2);
330    }
331
332    pub async fn test_multiple_variable(nvram: &mut dyn NvramStorage) {
333        let vendor1 = Guid::new_random();
334        let name1 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var1").as_bytes()).unwrap();
335        let vendor2 = Guid::new_random();
336        let name2 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var2").as_bytes()).unwrap();
337        let vendor3 = Guid::new_random();
338        let name3 = Ucs2LeSlice::from_slice_with_nul(wchz!(u16, "var3").as_bytes()).unwrap();
339        let attr = 0x1234;
340        let data = vec![0x1, 0x2, 0x3, 0x4, 0x5];
341        let timestamp = EFI_TIME::new_zeroed();
342
343        // add all variables to nvram
344        nvram
345            .set_variable(name1, vendor1, attr, data.clone(), timestamp)
346            .await
347            .unwrap();
348        nvram
349            .set_variable(name2, vendor2, attr, data.clone(), timestamp)
350            .await
351            .unwrap();
352        nvram
353            .set_variable(name3, vendor3, attr, data.clone(), timestamp)
354            .await
355            .unwrap();
356
357        let (result_attr, result_data, result_timestamp) =
358            nvram.get_variable(name1, vendor1).await.unwrap().unwrap();
359        assert_eq!(result_attr, attr);
360        assert_eq!(result_data, data);
361        assert_eq!(result_timestamp, timestamp);
362
363        let (result_attr, result_data, result_timestamp) =
364            nvram.get_variable(name2, vendor2).await.unwrap().unwrap();
365        assert_eq!(result_attr, attr);
366        assert_eq!(result_data, data);
367        assert_eq!(result_timestamp, timestamp);
368
369        let (result_attr, result_data, result_timestamp) =
370            nvram.get_variable(name3, vendor3).await.unwrap().unwrap();
371        assert_eq!(result_attr, attr);
372        assert_eq!(result_data, data);
373        assert_eq!(result_timestamp, timestamp);
374
375        // throw an append in there for good measure
376        let appended = nvram
377            .append_variable(name1, vendor1, vec![6, 7, 8], timestamp)
378            .await
379            .unwrap();
380        assert!(appended);
381
382        let (result_attr, result_data, result_timestamp) =
383            nvram.get_variable(name1, vendor1).await.unwrap().unwrap();
384        assert_eq!(result_attr, attr);
385        assert_eq!(result_data, (1..=8).collect::<Vec<u8>>());
386        assert_eq!(result_timestamp, timestamp);
387    }
388}
389
390#[cfg(test)]
391mod tests {
392    use super::impl_agnostic_tests;
393    use super::*;
394    use pal_async::async_test;
395
396    #[async_test]
397    async fn nvram_trait_single_variable() {
398        let mut nvram = InMemoryNvram::new();
399        impl_agnostic_tests::test_single_variable(&mut nvram).await;
400    }
401
402    #[async_test]
403    async fn nvram_trait_next() {
404        let mut nvram = InMemoryNvram::new();
405        impl_agnostic_tests::test_next(&mut nvram).await;
406    }
407
408    #[async_test]
409    async fn nvram_trait_multiple_variable() {
410        let mut nvram = InMemoryNvram::new();
411        impl_agnostic_tests::test_multiple_variable(&mut nvram).await;
412    }
413}
uefi_nvram_storage/in_memory.rs

uefi_nvram_storage/
in_memory.rs
