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
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 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 writer.write_u8(VARIANT_DICTIONARY_END)?;
186 Ok(())
187 }
188
189 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 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#[derive(Debug, PartialEq, Eq, Clone)]
215#[cfg_attr(feature = "serialization", derive(serde::Serialize))]
216#[non_exhaustive]
217pub enum VariantDictionaryValue {
218 UInt32(u32),
220
221 UInt64(u64),
223
224 Bool(bool),
226
227 Int32(i32),
229
230 Int64(i64),
232
233 String(String),
235
236 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#[derive(Debug, Error)]
347#[non_exhaustive]
348pub enum VariantDictionaryError {
349 #[error("Invalid variant dictionary version: {}", version)]
351 InvalidVersion {
352 version: u16,
354 },
355
356 #[error("Invalid value type: {}", value_type)]
358 InvalidValueType {
359 value_type: u8,
361 },
362
363 #[error("Missing key: {}", key)]
365 MissingKey {
366 key: String,
368 },
369
370 #[error("Mistyped value: {}", key)]
372 Mistyped {
373 key: String,
375 },
376
377 #[error("VariantDictionary did not end with null byte, when it should")]
379 NotTerminated,
380
381 #[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 let res = VariantDictionary::parse(&hex!("000104030000004142430400000015CD5B0700"))?;
407 assert_eq!(res.get_typed::<u32>("ABC")?, &123456789);
408
409 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}