bendy/decoding/
from_bencode.rs

1#[cfg(not(feature = "std"))]
2use alloc::{collections::BTreeMap, rc::Rc, string::String, vec::Vec};
3
4#[cfg(feature = "std")]
5use std::{
6    collections::{BTreeMap, HashMap},
7    hash::{BuildHasher, Hash},
8    rc::Rc,
9};
10
11use crate::{
12    decoding::{Decoder, Error, Object},
13    encoding::AsString,
14    state_tracker::StructureError,
15};
16
17/// Basic trait for bencode based value deserialization.
18pub trait FromBencode {
19    /// Maximum allowed depth of nested structures before the decoding should be aborted.
20    const EXPECTED_RECURSION_DEPTH: usize = 2048;
21
22    /// Deserialize an object from its byte representation.
23    fn from_bencode(bytes: &[u8]) -> Result<Self, Error>
24    where
25        Self: Sized,
26    {
27        let mut decoder = Decoder::new(bytes).with_max_depth(Self::EXPECTED_RECURSION_DEPTH);
28        let object = decoder.next_object()?;
29
30        object.map_or(
31            Err(Error::from(StructureError::UnexpectedEof)),
32            Self::decode_bencode_object,
33        )
34    }
35
36    /// Deserialize an object from its intermediate bencode representation.
37    fn decode_bencode_object(object: Object) -> Result<Self, Error>
38    where
39        Self: Sized;
40}
41
42macro_rules! impl_from_bencode_for_integer {
43    ($($type:ty)*) => {$(
44        impl FromBencode for $type {
45            const EXPECTED_RECURSION_DEPTH: usize = 0;
46
47            fn decode_bencode_object(object: Object) -> Result<Self, Error>
48            where
49                Self: Sized,
50            {
51                let content = object.try_into_integer()?;
52                let number = content.parse::<$type>()?;
53
54                Ok(number)
55            }
56        }
57    )*}
58}
59
60impl_from_bencode_for_integer!(u8 u16 u32 u64 u128 usize i8 i16 i32 i64 i128 isize);
61
62impl<ContentT: FromBencode> FromBencode for Vec<ContentT> {
63    const EXPECTED_RECURSION_DEPTH: usize = ContentT::EXPECTED_RECURSION_DEPTH + 1;
64
65    fn decode_bencode_object(object: Object) -> Result<Self, Error>
66    where
67        Self: Sized,
68    {
69        let mut list = object.try_into_list()?;
70        let mut results = Vec::new();
71
72        while let Some(object) = list.next_object()? {
73            let item = ContentT::decode_bencode_object(object)?;
74            results.push(item);
75        }
76
77        Ok(results)
78    }
79}
80
81impl FromBencode for String {
82    const EXPECTED_RECURSION_DEPTH: usize = 0;
83
84    fn decode_bencode_object(object: Object) -> Result<Self, Error>
85    where
86        Self: Sized,
87    {
88        let content = object.try_into_bytes()?;
89        let content = String::from_utf8(content.to_vec())?;
90
91        Ok(content)
92    }
93}
94
95impl<K, V> FromBencode for BTreeMap<K, V>
96where
97    K: FromBencode + Ord,
98    V: FromBencode,
99{
100    const EXPECTED_RECURSION_DEPTH: usize = V::EXPECTED_RECURSION_DEPTH + 1;
101
102    fn decode_bencode_object(object: Object) -> Result<Self, Error>
103    where
104        Self: Sized,
105    {
106        let mut dict = object.try_into_dictionary()?;
107        let mut result = BTreeMap::default();
108
109        while let Some((key, value)) = dict.next_pair()? {
110            let key = K::decode_bencode_object(Object::Bytes(key))?;
111            let value = V::decode_bencode_object(value)?;
112
113            result.insert(key, value);
114        }
115
116        Ok(result)
117    }
118}
119
120#[cfg(feature = "std")]
121impl<K, V, H> FromBencode for HashMap<K, V, H>
122where
123    K: FromBencode + Hash + Eq,
124    V: FromBencode,
125    H: BuildHasher + Default,
126{
127    const EXPECTED_RECURSION_DEPTH: usize = V::EXPECTED_RECURSION_DEPTH + 1;
128
129    fn decode_bencode_object(object: Object) -> Result<Self, Error>
130    where
131        Self: Sized,
132    {
133        let mut dict = object.try_into_dictionary()?;
134        let mut result = HashMap::default();
135
136        while let Some((key, value)) = dict.next_pair()? {
137            let key = K::decode_bencode_object(Object::Bytes(key))?;
138            let value = V::decode_bencode_object(value)?;
139
140            result.insert(key, value);
141        }
142
143        Ok(result)
144    }
145}
146
147impl<T: FromBencode> FromBencode for Rc<T> {
148    const EXPECTED_RECURSION_DEPTH: usize = T::EXPECTED_RECURSION_DEPTH;
149
150    fn decode_bencode_object(object: Object) -> Result<Self, Error>
151    where
152        Self: Sized,
153    {
154        T::decode_bencode_object(object).map(Rc::new)
155    }
156}
157
158impl FromBencode for AsString<Vec<u8>> {
159    const EXPECTED_RECURSION_DEPTH: usize = 0;
160
161    fn decode_bencode_object(object: Object) -> Result<Self, Error>
162    where
163        Self: Sized,
164    {
165        object.try_into_bytes().map(Vec::from).map(AsString)
166    }
167}
168
169#[cfg(test)]
170mod test {
171
172    #[cfg(not(feature = "std"))]
173    use alloc::{format, vec::Vec};
174
175    use crate::encoding::AsString;
176
177    use super::*;
178
179    #[test]
180    fn from_bencode_to_string_should_work_with_valid_input() {
181        let expected_message = "hello";
182        let serialized_message =
183            format!("{}:{}", expected_message.len(), expected_message).into_bytes();
184
185        let decoded_message = String::from_bencode(&serialized_message).unwrap();
186        assert_eq!(expected_message, decoded_message);
187    }
188
189    #[test]
190    fn from_bencode_to_as_string_should_work_with_valid_input() {
191        let expected_message = "hello";
192        let serialized_message =
193            format!("{}:{}", expected_message.len(), expected_message).into_bytes();
194
195        let decoded_vector = AsString::from_bencode(&serialized_message).unwrap();
196        assert_eq!(expected_message.as_bytes(), &decoded_vector.0[..]);
197    }
198
199    #[test]
200    #[should_panic(expected = "Num")]
201    fn from_bencode_to_as_string_should_fail_for_integer() {
202        AsString::<Vec<u8>>::from_bencode(&b"i1e"[..]).unwrap();
203    }
204
205    #[test]
206    #[should_panic(expected = "NestingTooDeep")]
207    fn from_bencode_to_as_string_should_fail_for_list() {
208        AsString::<Vec<u8>>::from_bencode(&b"l1:ae"[..]).unwrap();
209    }
210
211    #[test]
212    #[should_panic(expected = "NestingTooDeep")]
213    fn from_bencode_to_as_string_should_fail_for_dictionary() {
214        AsString::<Vec<u8>>::from_bencode(&b"d1:a1:ae"[..]).unwrap();
215    }
216}