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