Skip to main content

bencode_minimal/
value.rs

1use super::decoder::Decoder;
2use super::encoder::Encoder;
3use super::TryFromValue;
4use std::borrow::Cow;
5use std::collections::BTreeMap;
6
7/// An alias for [i64]
8pub type Int = i64;
9
10/// An alias for byte string ([Cow] of `[u8]`)
11pub type Str<'a> = Cow<'a, [u8]>;
12
13/// An alias for a list of [Value]s
14pub type List<'a> = Vec<Value<'a>>;
15
16/// An alias for a dictionary mapping byte strings to [Value]s
17pub type Dict<'a> = BTreeMap<Cow<'a, [u8]>, Value<'a>>;
18
19/// A Bencode value is either an [Int], a [Str], a [List] or a [Dict]
20///
21/// Note that [Value] carries a lifetime parameter for borrowed data. This is useful for
22/// two reasons:
23///
24/// 1. When encoding data, you can avoid unnecessary allocations by just referencing existing data.
25///    Usually, the object is serialized right away and the object is dropped afterwards.
26/// 2. When decoding data, all byte strings and dictionary keys are borrowed from the input buffer.
27///    This avoids unnecessary allocations and is key to performance. Especially since keys are
28///    often quite short, a separate [Vec] allocation per key would be very wasteful. The idea here is
29///    also that after decoding, the object is used and dispatched right away and then dropped.
30///    Actually required and useful data can be cloned out of the object if needed.
31///
32/// [Value]s are meant to be constructed using a combination of the provided macros [int!](super::int!),
33/// [str!](super::str!), [list!](super::list!) and [dict!](super::dict!):
34///
35/// ```rust
36/// use bencode_minimal::*;
37///
38/// let alice = "Alice".to_string();
39///
40/// let v = dict! {
41///     "age" => int!(42),
42///     "name" => str!("John"),                     // <-- static
43///     "friends" => list![
44///         str!(alice.as_str()),                   // <-- borrowed
45///         dict! {
46///             "name" => str!("Bob".to_string()),  // <-- owned
47///             "data" => str!(vec![48u8, 49, 50]), // <-- owned
48///         }
49///     ]
50/// };
51///
52/// let bin = v.encode();
53/// assert_eq!(&bin, b"d3:agei42e7:friendsl5:Aliced4:data3:0124:name3:Bobee4:name4:Johne");
54/// ```
55#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
56pub enum Value<'a> {
57    Int(Int),
58    Str(Str<'a>),
59    List(List<'a>),
60    Dict(Dict<'a>),
61}
62
63impl<'a> Value<'a> {
64    /// Assume the value is a dictionary and get the value for the given key, converted into the desired type
65    /// using [TryFromValue]
66    ///
67    /// Fails if the value is not a dictionary, the key does not exist or the value cannot be converted
68    /// into the desired type.
69    ///
70    /// ```rust
71    /// use bencode_minimal::*;
72    ///
73    /// let v = dict! {
74    ///     "t" => str!(b"1234"),
75    ///     "y" => str!(b"q"),
76    ///     "q" => str!(b"ping"),
77    ///     "a" => dict! {
78    ///         "id" => str!(vec![5u8; 20]),
79    ///     },
80    /// };
81    ///
82    /// let ping_query = || {
83    ///    let _ = v.get::<&str>("y").filter(|x| *x == "q")?;
84    ///    let _ = v.get::<&str>("q").filter(|x| *x == "ping")?;
85    ///    let a = v.get::<&Value>("a")?;
86    ///    let t = v.get::<&[u8]>("t")?;
87    ///    let id = a.get::<[u8;20]>("id")?;
88    ///    Some((t, id))
89    /// };
90    ///
91    /// let (t, id) = ping_query().unwrap();
92    /// assert_eq!(t, b"1234");
93    /// assert_eq!(id, [5u8; 20]);
94    /// ```
95    pub fn get<'b, T: TryFromValue<'b>>(&'b self, key: &'static str) -> Option<T> {
96        let x = self.try_into::<&'b Dict<'b>>()?;
97        let x = x.get(key.as_bytes())?;
98        x.try_into()
99    }
100
101    /// Try to convert the [Value] into the desired type using [TryFromValue]
102    ///
103    /// Fails if the value cannot be converted into the desired type.
104    pub fn try_into<'b, T: TryFromValue<'b>>(&'b self) -> Option<T> {
105        T::try_from(self)
106    }
107
108    /// Quick encoding into a [Vec]<[u8]>
109    ///
110    /// The returned vector is freshly allocated and has a capacity of 1500 bytes to
111    /// avoid multiple reallocations for typical use cases. Its length is adjusted to the
112    /// actual encoded size.
113    pub fn encode(&self) -> Vec<u8> {
114        let mut v = Vec::with_capacity(1500);
115        let mut e = Encoder::new(&mut v);
116        e.value(self);
117        v
118    }
119
120    /// Encode into the provided buffer
121    ///
122    /// The provided buffer is cleared before encoding and has the length of the encoded data afterwards.
123    /// Its capacity is increased as needed but never decreased. Make sure to reuse the buffer when possible
124    /// to avoid unnecessary allocations.Also, make sure to provide a buffer with sufficient initial capacity
125    /// to avoid multiple reallocations. When passing the same buffer multiple times, its capacity will grow
126    /// to the maximum size needed.
127    pub fn encode_into(&self, buf: &mut Vec<u8>) {
128        let mut e = Encoder::new(buf);
129        e.clear();
130        e.value(self);
131    }
132
133    /// Try to decode a [Value] from the provided buffer
134    ///
135    /// The `max_allocs` parameter limits the number of allocations that may be performed during decoding.
136    /// This is useful to avoid denial-of-service attacks by providing maliciously crafted input that would
137    /// cause excessive memory allocations. Each list item and dictionary entry counts as one allocation.
138    /// If the limit is exceeded, decoding fails and `None` is returned.
139    ///
140    /// The returned [Value] borrows all byte strings from the input buffer. The value can therefor not outlive
141    /// the input buffer. Either deconstruct the value right away (recommended) or use [Self::into_owned].
142    pub fn decode(buf: &'a [u8], max_allocs: usize) -> Option<Self> {
143        Decoder::new(buf, max_allocs).take_value()
144    }
145
146    /// Convert the value into an owned version
147    ///
148    /// All borrowed byte strings are cloned into owned [Vec]<[u8]>s. Byte strings that are already owned
149    /// are moved into the new value without cloning. All [Vec]s and [BTreeMap]s get unfortunately
150    /// recreated since there is no way to recycle them.
151    pub fn into_owned(self) -> Value<'static> {
152        match self {
153            Value::Int(i) => Value::Int(i),
154            Value::Str(s) => Value::Str(Cow::Owned(s.into_owned())),
155            Value::List(l) => Value::List(l.into_iter().map(Value::into_owned).collect()),
156            Value::Dict(d) => {
157                Value::Dict(d.into_iter().map(|(k, v)| (Cow::Owned(k.into_owned()), v.into_owned())).collect())
158            }
159        }
160    }
161}
162
163impl std::fmt::Debug for Value<'_> {
164    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
165        match self {
166            Value::Int(i) => write!(f, "{}", i),
167            Value::Str(s) => match std::str::from_utf8(s) {
168                Ok(s) => write!(f, "{:?}", s),
169                Err(_) => {
170                    for i in s.iter() {
171                        write!(f, "{:02x}", i)?;
172                    }
173                    Ok(())
174                }
175            },
176            Value::List(l) => f.debug_list().entries(l.iter()).finish(),
177            Value::Dict(d) => f
178                .debug_map()
179                .entries(d.iter().map(|(k, v)| {
180                    let k = match std::str::from_utf8(k) {
181                        Ok(s) => s.to_string(),
182                        Err(_) => format!("{:?}", k),
183                    };
184                    (k, v)
185                }))
186                .finish(),
187        }
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194
195    #[test]
196    fn test_int_01() {
197        let value = Value::Int(0);
198        let encoded = value.encode();
199        assert_eq!(&encoded, b"i0e");
200        let value_ = Value::decode(&encoded, 0);
201        assert_eq!(value_, Some(value));
202    }
203
204    #[test]
205    fn test_int_02() {
206        let value = Value::Int(1);
207        let encoded = value.encode();
208        assert_eq!(&encoded, b"i1e");
209        let value_ = Value::decode(&encoded, 0);
210        assert_eq!(value_, Some(value));
211    }
212
213    #[test]
214    fn test_int_03() {
215        let value = Value::Int(-1);
216        let encoded = value.encode();
217        assert_eq!(&encoded, b"i-1e");
218        let value_ = Value::decode(&encoded, 0);
219        assert_eq!(value_, Some(value));
220    }
221
222    #[test]
223    fn test_int_04() {
224        let value = Value::Int(10);
225        let encoded = value.encode();
226        assert_eq!(&encoded, b"i10e");
227        let value_ = Value::decode(&encoded, 0);
228        assert_eq!(value_, Some(value));
229    }
230
231    #[test]
232    fn test_int_05() {
233        let value = Value::Int(-10);
234        let encoded = value.encode();
235        assert_eq!(&encoded, b"i-10e");
236        let value_ = Value::decode(&encoded, 0);
237        assert_eq!(value_, Some(value));
238    }
239
240    #[test]
241    fn test_int_06() {
242        let value = Value::Int(42);
243        let encoded = value.encode();
244        assert_eq!(&encoded, b"i42e");
245        let value_ = Value::decode(&encoded, 0);
246        assert_eq!(value_, Some(value));
247    }
248
249    #[test]
250    fn test_int_07() {
251        let value = Value::Int(-42);
252        let encoded = value.encode();
253        assert_eq!(&encoded, b"i-42e");
254        let value_ = Value::decode(&encoded, 0);
255        assert_eq!(value_, Some(value));
256    }
257
258    #[test]
259    fn test_str_01() {
260        let value = Value::Str(Cow::Borrowed(b""));
261        let encoded = value.encode();
262        assert_eq!(&encoded, b"0:");
263        let value_ = Value::decode(&encoded, 0);
264        assert_eq!(value_, Some(value));
265    }
266
267    #[test]
268    fn test_str_02() {
269        let value = Value::Str(Cow::Borrowed(b":"));
270        let encoded = value.encode();
271        assert_eq!(&encoded, b"1::");
272        let value_ = Value::decode(&encoded, 0);
273        assert_eq!(value_, Some(value));
274    }
275
276    #[test]
277    fn test_str_03() {
278        let value = Value::Str(Cow::Borrowed(b"hello"));
279        let encoded = value.encode();
280        assert_eq!(&encoded, b"5:hello");
281        let value_ = Value::decode(&encoded, 0);
282        assert_eq!(value_, Some(value));
283    }
284
285    #[test]
286    fn test_str_04() {
287        let value = Value::Str(Cow::Borrowed(b"helloworld"));
288        let encoded = value.encode();
289        assert_eq!(&encoded, b"10:helloworld");
290        let value_ = Value::decode(&encoded, 0);
291        assert_eq!(value_, Some(value));
292    }
293
294    #[test]
295    fn test_list_01() {
296        let value = Value::List(vec![]);
297        let encoded = value.encode();
298        assert_eq!(&encoded, b"le");
299        let value_ = Value::decode(&encoded, 10);
300        assert_eq!(value_, Some(value));
301    }
302
303    #[test]
304    fn test_list_02() {
305        let value = Value::List(vec![Value::Int(42), Value::Str(Cow::Borrowed(b"hello"))]);
306        let encoded = value.encode();
307        assert_eq!(&encoded, b"li42e5:helloe");
308        let value_ = Value::decode(&encoded, 10);
309        assert_eq!(value_, Some(value));
310    }
311
312    #[test]
313    fn test_dict_01() {
314        let value = Value::Dict(BTreeMap::new());
315        let encoded = value.encode();
316        assert_eq!(&encoded, b"de");
317        let value_ = Value::decode(&encoded, 10);
318        assert_eq!(value_, Some(value));
319    }
320
321    #[test]
322    fn test_dict_02() {
323        let mut dict = BTreeMap::new();
324        dict.insert(b"age".into(), Value::Int(42));
325        dict.insert(b"name".into(), Value::Str(Cow::Borrowed(b"John")));
326        let value = Value::Dict(dict);
327        let encoded = value.encode();
328        assert_eq!(&encoded, b"d3:agei42e4:name4:Johne");
329        let value_ = Value::decode(&encoded, 10);
330        assert_eq!(value_, Some(value));
331    }
332
333    #[test]
334    fn test_dict_03_reversed_order() {
335        let mut dict = BTreeMap::new();
336        dict.insert(b"name".into(), Value::Str(Cow::Borrowed(b"John")));
337        dict.insert(b"age".into(), Value::Int(42));
338        let value = Value::Dict(dict);
339
340        let encoded = b"d4:name4:John3:agei42ee";
341        let value_ = Value::decode(encoded.as_ref(), 10);
342        assert_eq!(value_, Some(value));
343    }
344
345    #[test]
346    fn test_dict_04_duplicate_keys() {
347        let encoded = b"d3:agei30e3:agei40ee";
348        let value = Value::decode(encoded.as_ref(), 10);
349        assert!(value.is_none());
350    }
351
352    #[test]
353    fn test_max_alloc_int() {
354        let encoded = b"i42e";
355        let value = Value::decode(encoded.as_ref(), 0);
356        assert!(value.is_some());
357    }
358
359    #[test]
360    fn test_max_alloc_str() {
361        let encoded = b"5:hello";
362        let value_ = Value::decode(encoded.as_ref(), 0);
363        assert!(value_.is_some());
364    }
365
366    #[test]
367    fn test_max_alloc_list_empty() {
368        let encoded = b"le";
369        let value = Value::decode(encoded.as_ref(), 0);
370        assert!(value.is_some());
371    }
372
373    #[test]
374    fn test_max_alloc_list_one() {
375        let encoded = b"li42ee";
376
377        let value = Value::decode(encoded.as_ref(), 0);
378        assert!(value.is_none());
379
380        let value = Value::decode(encoded.as_ref(), 1);
381        assert!(value.is_some());
382    }
383
384    #[test]
385    fn test_max_alloc_list_two() {
386        let encoded = b"li1ei2ee";
387
388        let value = Value::decode(encoded.as_ref(), 1);
389        assert!(value.is_none());
390
391        let value = Value::decode(encoded.as_ref(), 2);
392        assert!(value.is_some());
393    }
394
395    #[test]
396    fn test_max_alloc_dict_empty() {
397        let encoded = b"de";
398        let value = Value::decode(encoded.as_ref(), 0);
399        assert!(value.is_some());
400    }
401
402    #[test]
403    fn test_max_alloc_dict_one() {
404        let encoded = b"d3:agei42ee";
405
406        let value = Value::decode(encoded.as_ref(), 0);
407        assert!(value.is_none());
408
409        let value = Value::decode(encoded.as_ref(), 1);
410        assert!(value.is_some());
411    }
412
413    #[test]
414    fn test_max_alloc_dict_two() {
415        let encoded = b"d3:agei42e4:name4:Johne";
416
417        let value = Value::decode(encoded.as_ref(), 1);
418        assert!(value.is_none());
419
420        let value = Value::decode(encoded.as_ref(), 2);
421        assert!(value.is_some());
422    }
423}