Skip to main content

keepass/format/
variant_dictionary.rs

1#[cfg(feature = "save_kdbx4")]
2use byteorder::WriteBytesExt;
3use byteorder::{ByteOrder, LittleEndian};
4#[cfg(feature = "save_kdbx4")]
5use std::io::Write;
6use std::{
7    collections::HashMap,
8    ops::{Deref, DerefMut},
9};
10use thiserror::Error;
11
12#[cfg(feature = "save_kdbx4")]
13use crate::format::io::WriteLengthTaggedExt;
14
15pub const VARIANT_DICTIONARY_VERSION: u16 = 0x100;
16pub const VARIANT_DICTIONARY_END: u8 = 0x0;
17
18pub const U32_TYPE_ID: u8 = 0x04;
19pub const U64_TYPE_ID: u8 = 0x05;
20pub const BOOL_TYPE_ID: u8 = 0x08;
21pub const I32_TYPE_ID: u8 = 0x0c;
22pub const I64_TYPE_ID: u8 = 0x0d;
23pub const STR_TYPE_ID: u8 = 0x18;
24pub const BYTES_TYPE_ID: u8 = 0x42;
25
26/// A dictionary of key-value pairs, with typed values
27#[derive(Debug, PartialEq, Eq, Clone, Default)]
28#[cfg_attr(feature = "serialization", derive(serde::Serialize))]
29pub struct VariantDictionary(HashMap<String, VariantDictionaryValue>);
30
31impl Deref for VariantDictionary {
32    type Target = HashMap<String, VariantDictionaryValue>;
33
34    fn deref(&self) -> &Self::Target {
35        &self.0
36    }
37}
38
39impl DerefMut for VariantDictionary {
40    fn deref_mut(&mut self) -> &mut Self::Target {
41        &mut self.0
42    }
43}
44
45impl VariantDictionary {
46    /// Create a new, empty VariantDictionary
47    pub fn new() -> Self {
48        Self(HashMap::new())
49    }
50
51    pub(crate) fn parse(buffer: &[u8]) -> Result<VariantDictionary, VariantDictionaryError> {
52        let version = buffer.get(0..2).ok_or(VariantDictionaryError::UnexpectedEof)?;
53        let version = LittleEndian::read_u16(version);
54
55        if version != VARIANT_DICTIONARY_VERSION {
56            return Err(VariantDictionaryError::InvalidVersion { version });
57        }
58
59        let mut pos = 2;
60        let mut data = HashMap::new();
61
62        while pos + 9 < buffer.len() {
63            let value_type = *buffer.get(pos).ok_or(VariantDictionaryError::UnexpectedEof)?;
64            pos += 1;
65
66            let key_length = buffer
67                .get(pos..(pos + 4))
68                .ok_or(VariantDictionaryError::UnexpectedEof)?;
69            let key_length = LittleEndian::read_u32(key_length) as usize;
70            pos += 4;
71
72            let key = buffer
73                .get(pos..(pos + key_length))
74                .ok_or(VariantDictionaryError::UnexpectedEof)?;
75            let key = String::from_utf8_lossy(key).to_string();
76            pos += key_length;
77
78            let value_length = buffer
79                .get(pos..(pos + 4))
80                .ok_or(VariantDictionaryError::UnexpectedEof)?;
81            let value_length = LittleEndian::read_u32(value_length) as usize;
82            pos += 4;
83
84            let value_buffer = buffer
85                .get(pos..(pos + value_length))
86                .ok_or(VariantDictionaryError::UnexpectedEof)?;
87            pos += value_length;
88
89            let value = match value_type {
90                U32_TYPE_ID => VariantDictionaryValue::UInt32(LittleEndian::read_u32(
91                    value_buffer.get(0..4).ok_or(VariantDictionaryError::UnexpectedEof)?,
92                )),
93                U64_TYPE_ID => VariantDictionaryValue::UInt64(LittleEndian::read_u64(
94                    value_buffer.get(0..8).ok_or(VariantDictionaryError::UnexpectedEof)?,
95                )),
96                BOOL_TYPE_ID => VariantDictionaryValue::Bool(value_buffer != [0]),
97                I32_TYPE_ID => VariantDictionaryValue::Int32(LittleEndian::read_i32(
98                    value_buffer.get(0..4).ok_or(VariantDictionaryError::UnexpectedEof)?,
99                )),
100                I64_TYPE_ID => VariantDictionaryValue::Int64(LittleEndian::read_i64(
101                    value_buffer.get(0..8).ok_or(VariantDictionaryError::UnexpectedEof)?,
102                )),
103                STR_TYPE_ID => {
104                    VariantDictionaryValue::String(String::from_utf8_lossy(value_buffer).to_string())
105                }
106                BYTES_TYPE_ID => VariantDictionaryValue::ByteArray(value_buffer.to_vec()),
107                _ => {
108                    return Err(VariantDictionaryError::InvalidValueType { value_type });
109                }
110            };
111
112            data.insert(key, value);
113        }
114
115        if pos == buffer.len()
116            || *buffer.get(pos).ok_or(VariantDictionaryError::UnexpectedEof)? != VARIANT_DICTIONARY_END
117        {
118            // even though we can determine when to stop parsing a VariantDictionary by where we
119            // are in the buffer, there should always be a value_type = 0 entry to denote that a
120            // VariantDictionary is finished
121            return Err(VariantDictionaryError::NotTerminated);
122        }
123
124        Ok(VariantDictionary(data))
125    }
126
127    #[cfg(feature = "save_kdbx4")]
128    pub(crate) fn dump(&self, writer: &mut dyn Write) -> Result<(), std::io::Error> {
129        writer.write_u16::<LittleEndian>(VARIANT_DICTIONARY_VERSION)?;
130
131        for (field_name, field_value) in &self.0 {
132            match field_value {
133                VariantDictionaryValue::UInt32(value) => {
134                    writer.write_u8(U32_TYPE_ID)?;
135                    writer.write_with_len(field_name.as_bytes())?;
136                    writer.write_u32::<LittleEndian>(4)?;
137                    writer.write_u32::<LittleEndian>(*value)?;
138                }
139                VariantDictionaryValue::UInt64(value) => {
140                    writer.write_u8(U64_TYPE_ID)?;
141                    writer.write_with_len(field_name.as_bytes())?;
142                    writer.write_u32::<LittleEndian>(8)?;
143                    writer.write_u64::<LittleEndian>(*value)?;
144                }
145                VariantDictionaryValue::Bool(value) => {
146                    writer.write_u8(BOOL_TYPE_ID)?;
147                    writer.write_with_len(field_name.as_bytes())?;
148                    writer.write_u32::<LittleEndian>(1)?;
149                    writer.write_u8(if *value { 1 } else { 0 })?;
150                }
151                VariantDictionaryValue::Int32(value) => {
152                    writer.write_u8(I32_TYPE_ID)?;
153                    writer.write_with_len(field_name.as_bytes())?;
154                    writer.write_u32::<LittleEndian>(4)?;
155                    writer.write_i32::<LittleEndian>(*value)?;
156                }
157                VariantDictionaryValue::Int64(value) => {
158                    writer.write_u8(I64_TYPE_ID)?;
159                    writer.write_with_len(field_name.as_bytes())?;
160                    writer.write_u32::<LittleEndian>(8)?;
161                    writer.write_i64::<LittleEndian>(*value)?;
162                }
163                VariantDictionaryValue::String(value) => {
164                    writer.write_u8(STR_TYPE_ID)?;
165                    writer.write_with_len(field_name.as_bytes())?;
166                    writer.write_with_len(value.as_bytes())?;
167                }
168                VariantDictionaryValue::ByteArray(value) => {
169                    writer.write_u8(BYTES_TYPE_ID)?;
170                    writer.write_with_len(field_name.as_bytes())?;
171                    writer.write_with_len(value)?;
172                }
173            };
174        }
175
176        // signify end of variant dictionary
177        writer.write_u8(VARIANT_DICTIONARY_END)?;
178        Ok(())
179    }
180
181    /// Get a value from the VariantDictionary, returning an error if the key is missing or the
182    /// value is of the wrong type
183    pub fn get_typed<'a, T: 'a>(&'a self, key: &str) -> Result<&'a T, VariantDictionaryError>
184    where
185        &'a VariantDictionaryValue: Into<Option<&'a T>>,
186    {
187        let vdv = self
188            .0
189            .get(key)
190            .ok_or_else(|| VariantDictionaryError::MissingKey { key: key.to_owned() })?;
191
192        vdv.into()
193            .ok_or_else(|| VariantDictionaryError::Mistyped { key: key.to_owned() })
194    }
195
196    /// Set a value in the VariantDictionary
197    pub fn set<T>(&mut self, key: &str, value: T)
198    where
199        T: Into<VariantDictionaryValue>,
200    {
201        self.insert(key.to_string(), value.into());
202    }
203}
204
205/// A value in a VariantDictionary, which can be one of several types
206#[derive(Debug, PartialEq, Eq, Clone)]
207#[cfg_attr(feature = "serialization", derive(serde::Serialize))]
208#[non_exhaustive]
209pub enum VariantDictionaryValue {
210    /// An unsigned 32-bit integer value
211    UInt32(u32),
212
213    /// An unsigned 64-bit integer value
214    UInt64(u64),
215
216    /// A boolean value
217    Bool(bool),
218
219    /// A signed 32-bit integer value
220    Int32(i32),
221
222    /// A signed 64-bit integer value
223    Int64(i64),
224
225    /// A UTF-8 encoded string value
226    String(String),
227
228    /// A byte array value
229    ByteArray(Vec<u8>),
230}
231
232impl From<u32> for VariantDictionaryValue {
233    fn from(v: u32) -> Self {
234        VariantDictionaryValue::UInt32(v)
235    }
236}
237
238impl From<u64> for VariantDictionaryValue {
239    fn from(v: u64) -> Self {
240        VariantDictionaryValue::UInt64(v)
241    }
242}
243
244impl From<i32> for VariantDictionaryValue {
245    fn from(v: i32) -> Self {
246        VariantDictionaryValue::Int32(v)
247    }
248}
249
250impl From<i64> for VariantDictionaryValue {
251    fn from(v: i64) -> Self {
252        VariantDictionaryValue::Int64(v)
253    }
254}
255
256impl From<bool> for VariantDictionaryValue {
257    fn from(v: bool) -> Self {
258        VariantDictionaryValue::Bool(v)
259    }
260}
261
262impl From<String> for VariantDictionaryValue {
263    fn from(v: String) -> Self {
264        VariantDictionaryValue::String(v)
265    }
266}
267
268impl From<Vec<u8>> for VariantDictionaryValue {
269    fn from(v: Vec<u8>) -> Self {
270        VariantDictionaryValue::ByteArray(v)
271    }
272}
273
274impl<'a> From<&'a VariantDictionaryValue> for Option<&'a u32> {
275    fn from(val: &'a VariantDictionaryValue) -> Self {
276        match val {
277            VariantDictionaryValue::UInt32(v) => Some(v),
278            _ => None,
279        }
280    }
281}
282
283impl<'a> From<&'a VariantDictionaryValue> for Option<&'a u64> {
284    fn from(val: &'a VariantDictionaryValue) -> Self {
285        match val {
286            VariantDictionaryValue::UInt64(v) => Some(v),
287            _ => None,
288        }
289    }
290}
291
292impl<'a> From<&'a VariantDictionaryValue> for Option<&'a bool> {
293    fn from(val: &'a VariantDictionaryValue) -> Self {
294        match val {
295            VariantDictionaryValue::Bool(v) => Some(v),
296            _ => None,
297        }
298    }
299}
300
301impl<'a> From<&'a VariantDictionaryValue> for Option<&'a i32> {
302    fn from(val: &'a VariantDictionaryValue) -> Self {
303        match val {
304            VariantDictionaryValue::Int32(v) => Some(v),
305            _ => None,
306        }
307    }
308}
309
310impl<'a> From<&'a VariantDictionaryValue> for Option<&'a i64> {
311    fn from(val: &'a VariantDictionaryValue) -> Self {
312        match val {
313            VariantDictionaryValue::Int64(v) => Some(v),
314            _ => None,
315        }
316    }
317}
318
319impl<'a> From<&'a VariantDictionaryValue> for Option<&'a String> {
320    fn from(val: &'a VariantDictionaryValue) -> Self {
321        match val {
322            VariantDictionaryValue::String(v) => Some(v),
323            _ => None,
324        }
325    }
326}
327
328impl<'a> From<&'a VariantDictionaryValue> for Option<&'a Vec<u8>> {
329    fn from(val: &'a VariantDictionaryValue) -> Self {
330        match val {
331            VariantDictionaryValue::ByteArray(v) => Some(v),
332            _ => None,
333        }
334    }
335}
336
337/// Errors while parsing a VariantDictionary
338#[derive(Debug, Error)]
339#[non_exhaustive]
340pub enum VariantDictionaryError {
341    /// An invalid VariantDictionary version was encountered.
342    #[error("Invalid variant dictionary version: {}", version)]
343    InvalidVersion {
344        /// The version number that was encountered
345        version: u16,
346    },
347
348    /// An invalid value type was encountered while parsing a VariantDictionary.
349    #[error("Invalid value type: {}", value_type)]
350    InvalidValueType {
351        /// The value type identifier that was encountered
352        value_type: u8,
353    },
354
355    /// A required key was missing from the VariantDictionary.
356    #[error("Missing key: {}", key)]
357    MissingKey {
358        /// The name of the missing key
359        key: String,
360    },
361
362    /// A value was found for the specified key, but it was of an unexpected type.
363    #[error("Mistyped value: {}", key)]
364    Mistyped {
365        /// The name of the key whose value was mistyped
366        key: String,
367    },
368
369    /// The VariantDictionary did not end with a null byte (0x00) as expected.
370    #[error("VariantDictionary did not end with null byte, when it should")]
371    NotTerminated,
372
373    /// An unexpected end of file was encountered while parsing the VariantDictionary
374    #[error("Unexpected end of file while parsing VariantDictionary")]
375    UnexpectedEof,
376}
377
378#[allow(clippy::unwrap_used)]
379#[cfg(test)]
380mod variant_dictionary_tests {
381    use hex_literal::hex;
382
383    use super::*;
384
385    #[test]
386    fn parsing_errors() -> Result<(), VariantDictionaryError> {
387        let res = VariantDictionary::parse("not-a-variant-dictionary".as_bytes());
388        assert!(matches!(res, Err(VariantDictionaryError::InvalidVersion { .. })));
389
390        let res = VariantDictionary::parse(&hex!("0001"));
391        assert!(matches!(res, Err(VariantDictionaryError::NotTerminated)));
392
393        let res = VariantDictionary::parse(&hex!("000100"));
394        assert!(res.is_ok());
395
396        //                                        ver t key_len key   val_len value   termination
397        //                                        |   | |       |     |       |       |
398        let res = VariantDictionary::parse(&hex!("000104030000004142430400000015CD5B0700"))?;
399        assert_eq!(res.get_typed::<u32>("ABC")?, &123456789);
400
401        //                                        ver t key_len key val_len termination
402        //                                        |   | |       |   |       |
403        let res = VariantDictionary::parse(&hex!("0001AA0200000041420000000000"));
404        dbg!(&res);
405        assert!(matches!(
406            res,
407            Err(VariantDictionaryError::InvalidValueType { value_type: 0xAA })
408        ));
409
410        Ok(())
411    }
412
413    #[test]
414    #[cfg(feature = "save_kdbx4")]
415    fn variant_dictionary() {
416        let mut vd = VariantDictionary::new();
417
418        vd.set("a-u32", 42u32);
419        vd.set("a-u64", 1337u64);
420        vd.set("a-i32", -2i32);
421        vd.set("a-i64", -31337i64);
422        vd.set("a-bool", true);
423        vd.set("a-string", "Testing".to_string());
424        vd.set("a-bytes", "testing".as_bytes().to_vec());
425
426        assert!(vd.get_typed::<bool>("key-not-exist").is_err());
427
428        assert!(vd.get_typed::<u32>("a-string").is_err());
429        assert!(vd.get_typed::<u64>("a-string").is_err());
430        assert!(vd.get_typed::<i32>("a-string").is_err());
431        assert!(vd.get_typed::<i64>("a-string").is_err());
432        assert!(vd.get_typed::<bool>("a-string").is_err());
433        assert!(vd.get_typed::<String>("a-bytes").is_err());
434        assert!(vd.get_typed::<Vec<u8>>("a-string").is_err());
435
436        assert_eq!(vd.get_typed::<u32>("a-u32").unwrap(), &42u32);
437        assert_eq!(vd.get_typed::<u64>("a-u64").unwrap(), &1337u64);
438        assert_eq!(vd.get_typed::<i32>("a-i32").unwrap(), &-2i32);
439        assert_eq!(vd.get_typed::<i64>("a-i64").unwrap(), &-31337i64);
440        assert_eq!(vd.get_typed::<bool>("a-bool").unwrap(), &true);
441        assert_eq!(vd.get_typed::<String>("a-string").unwrap(), "Testing");
442        assert_eq!(vd.get_typed::<Vec<u8>>("a-bytes").unwrap(), "testing".as_bytes());
443
444        let mut vd_data = Vec::new();
445        vd.dump(&mut vd_data).unwrap();
446
447        let vd_parsed = VariantDictionary::parse(&vd_data).unwrap();
448        assert_eq!(vd_parsed, vd);
449    }
450}