bendy/encoding/
to_bencode.rs

1#[cfg(not(feature = "std"))]
2use alloc::{
3    collections::{BTreeMap, LinkedList, VecDeque},
4    rc::Rc,
5    string::String,
6    sync::Arc,
7    vec::Vec,
8};
9
10#[cfg(feature = "std")]
11use std::{
12    collections::{BTreeMap, HashMap, LinkedList, VecDeque},
13    hash::{BuildHasher, Hash},
14    rc::Rc,
15    sync::Arc,
16};
17
18use crate::encoding::{Encoder, Error, SingleItemEncoder};
19
20/// An object that can be encoded into a single bencode object
21pub trait ToBencode {
22    /// The maximum depth that this object could encode to. Leaves do not consume a level, so an
23    /// `i1e` has depth 0 and `li1ee` has depth 1.
24    const MAX_DEPTH: usize;
25
26    /// Encode this object into the bencode stream
27    fn encode(&self, encoder: SingleItemEncoder) -> Result<(), Error>;
28
29    /// Encode this object to a byte string
30    fn to_bencode(&self) -> Result<Vec<u8>, Error> {
31        let mut encoder = Encoder::new().with_max_depth(Self::MAX_DEPTH);
32        encoder.emit_with(|e| self.encode(e).map_err(Error::into))?;
33
34        let bytes = encoder.get_output()?;
35        Ok(bytes)
36    }
37}
38
39/// Wrapper to allow `Vec<u8>` encoding as bencode string element.
40#[derive(Clone, Copy, Debug, Default, Hash, Eq, PartialEq, PartialOrd, Ord)]
41pub struct AsString<I>(pub I);
42
43// Forwarding impls
44impl<'a, E: 'a + ToBencode + Sized> ToBencode for &'a E {
45    const MAX_DEPTH: usize = E::MAX_DEPTH;
46
47    fn encode(&self, encoder: SingleItemEncoder) -> Result<(), Error> {
48        E::encode(self, encoder)
49    }
50}
51
52#[cfg(feature = "std")]
53impl<E: ToBencode> ToBencode for Box<E> {
54    const MAX_DEPTH: usize = E::MAX_DEPTH;
55
56    fn encode(&self, encoder: SingleItemEncoder) -> Result<(), Error> {
57        E::encode(&*self, encoder)
58    }
59}
60
61impl<E: ToBencode> ToBencode for Rc<E> {
62    const MAX_DEPTH: usize = E::MAX_DEPTH;
63
64    fn encode(&self, encoder: SingleItemEncoder) -> Result<(), Error> {
65        E::encode(&*self, encoder)
66    }
67}
68
69impl<E: ToBencode> ToBencode for Arc<E> {
70    const MAX_DEPTH: usize = E::MAX_DEPTH;
71
72    fn encode(&self, encoder: SingleItemEncoder) -> Result<(), Error> {
73        E::encode(&*self, encoder)
74    }
75}
76
77// Base type impls
78impl<'a> ToBencode for &'a str {
79    const MAX_DEPTH: usize = 0;
80
81    fn encode(&self, encoder: SingleItemEncoder) -> Result<(), Error> {
82        encoder.emit_str(self).map_err(Error::from)
83    }
84}
85
86impl ToBencode for String {
87    const MAX_DEPTH: usize = 0;
88
89    fn encode(&self, encoder: SingleItemEncoder) -> Result<(), Error> {
90        encoder.emit_str(self).map_err(Error::from)
91    }
92}
93
94macro_rules! impl_encodable_integer {
95    ($($type:ty)*) => {$(
96        impl ToBencode for $type {
97            const MAX_DEPTH: usize = 1;
98
99            fn encode(&self, encoder: SingleItemEncoder) -> Result<(), Error> {
100                encoder.emit_int(*self).map_err(Error::from)
101            }
102        }
103    )*}
104}
105
106impl_encodable_integer!(u8 u16 u32 u64 u128 usize i8 i16 i32 i64 i128 isize);
107
108macro_rules! impl_encodable_iterable {
109    ($($type:ident)*) => {$(
110        impl <ContentT> ToBencode for $type<ContentT>
111        where
112            ContentT: ToBencode
113        {
114            const MAX_DEPTH: usize = ContentT::MAX_DEPTH + 1;
115
116            fn encode(&self, encoder: SingleItemEncoder) -> Result<(), Error> {
117                encoder.emit_list(|e| {
118                    for item in self {
119                        e.emit(item)?;
120                    }
121                    Ok(())
122                })?;
123
124                Ok(())
125            }
126        }
127    )*}
128}
129
130impl_encodable_iterable!(Vec VecDeque LinkedList);
131
132impl<'a, ContentT> ToBencode for &'a [ContentT]
133where
134    ContentT: ToBencode,
135{
136    const MAX_DEPTH: usize = ContentT::MAX_DEPTH + 1;
137
138    fn encode(&self, encoder: SingleItemEncoder) -> Result<(), Error> {
139        encoder.emit_list(|e| {
140            for item in *self {
141                e.emit(item)?;
142            }
143            Ok(())
144        })?;
145
146        Ok(())
147    }
148}
149
150impl<K: AsRef<[u8]>, V: ToBencode> ToBencode for BTreeMap<K, V> {
151    const MAX_DEPTH: usize = V::MAX_DEPTH + 1;
152
153    fn encode(&self, encoder: SingleItemEncoder) -> Result<(), Error> {
154        encoder.emit_dict(|mut e| {
155            for (k, v) in self {
156                e.emit_pair(k.as_ref(), v)?;
157            }
158            Ok(())
159        })?;
160
161        Ok(())
162    }
163}
164
165#[cfg(feature = "std")]
166impl<K, V, S> ToBencode for HashMap<K, V, S>
167where
168    K: AsRef<[u8]> + Eq + Hash,
169    V: ToBencode,
170    S: BuildHasher,
171{
172    const MAX_DEPTH: usize = V::MAX_DEPTH + 1;
173
174    fn encode(&self, encoder: SingleItemEncoder) -> Result<(), Error> {
175        encoder.emit_dict(|mut e| {
176            let mut pairs = self
177                .iter()
178                .map(|(k, v)| (k.as_ref(), v))
179                .collect::<Vec<_>>();
180            pairs.sort_by_key(|&(k, _)| k);
181            for (k, v) in pairs {
182                e.emit_pair(k, v)?;
183            }
184            Ok(())
185        })?;
186
187        Ok(())
188    }
189}
190
191impl<I> ToBencode for AsString<I>
192where
193    I: AsRef<[u8]>,
194{
195    const MAX_DEPTH: usize = 1;
196
197    fn encode(&self, encoder: SingleItemEncoder) -> Result<(), Error> {
198        encoder.emit_bytes(self.0.as_ref())?;
199        Ok(())
200    }
201}
202
203impl<I> AsRef<[u8]> for AsString<I>
204where
205    I: AsRef<[u8]>,
206{
207    fn as_ref(&self) -> &'_ [u8] {
208        self.0.as_ref()
209    }
210}
211
212impl<'a, I> From<&'a [u8]> for AsString<I>
213where
214    I: From<&'a [u8]>,
215{
216    fn from(content: &'a [u8]) -> Self {
217        AsString(I::from(content))
218    }
219}
220
221#[cfg(test)]
222mod test {
223
224    #[cfg(not(feature = "std"))]
225    use alloc::{borrow::ToOwned, vec};
226
227    use super::*;
228
229    struct Foo {
230        bar: u32,
231        baz: Vec<String>,
232        qux: Vec<u8>,
233    }
234
235    impl ToBencode for Foo {
236        const MAX_DEPTH: usize = 2;
237
238        fn encode(&self, encoder: SingleItemEncoder) -> Result<(), Error> {
239            encoder.emit_dict(|mut e| {
240                e.emit_pair(b"bar", &self.bar)?;
241                e.emit_pair(b"baz", &self.baz)?;
242                e.emit_pair(b"qux", AsString(&self.qux))?;
243                Ok(())
244            })?;
245
246            Ok(())
247        }
248    }
249
250    #[test]
251    fn simple_encodable_works() {
252        let mut encoder = Encoder::new();
253        encoder
254            .emit(Foo {
255                bar: 5,
256                baz: vec!["foo".to_owned(), "bar".to_owned()],
257                qux: b"qux".to_vec(),
258            })
259            .unwrap();
260        assert_eq!(
261            &encoder.get_output().unwrap()[..],
262            &b"d3:bari5e3:bazl3:foo3:bare3:qux3:quxe"[..]
263        );
264    }
265}