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#[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 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 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 writer.write_u8(VARIANT_DICTIONARY_END)?;
178 Ok(())
179 }
180
181 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 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#[derive(Debug, PartialEq, Eq, Clone)]
207#[cfg_attr(feature = "serialization", derive(serde::Serialize))]
208#[non_exhaustive]
209pub enum VariantDictionaryValue {
210 UInt32(u32),
212
213 UInt64(u64),
215
216 Bool(bool),
218
219 Int32(i32),
221
222 Int64(i64),
224
225 String(String),
227
228 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#[derive(Debug, Error)]
339#[non_exhaustive]
340pub enum VariantDictionaryError {
341 #[error("Invalid variant dictionary version: {}", version)]
343 InvalidVersion {
344 version: u16,
346 },
347
348 #[error("Invalid value type: {}", value_type)]
350 InvalidValueType {
351 value_type: u8,
353 },
354
355 #[error("Missing key: {}", key)]
357 MissingKey {
358 key: String,
360 },
361
362 #[error("Mistyped value: {}", key)]
364 Mistyped {
365 key: String,
367 },
368
369 #[error("VariantDictionary did not end with null byte, when it should")]
371 NotTerminated,
372
373 #[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 let res = VariantDictionary::parse(&hex!("000104030000004142430400000015CD5B0700"))?;
399 assert_eq!(res.get_typed::<u32>("ABC")?, &123456789);
400
401 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}