Skip to main content

librqbit_bencode/
serialize.rs

1use serde::{Serialize, Serializer, ser::Impossible};
2
3use buffers::ByteBufOwned;
4
5#[derive(thiserror::Error, Debug)]
6pub enum Error {
7    #[error(transparent)]
8    Io(std::io::Error),
9    #[error("{0}")]
10    Custom(Box<Box<str>>), // double box to reduce size
11    #[error("{0}")]
12    Text(&'static &'static str),
13}
14
15impl serde::ser::Error for Error {
16    fn custom<T>(msg: T) -> Self
17    where
18        T: std::fmt::Display,
19    {
20        Error::Custom(Box::new(msg.to_string().into_boxed_str()))
21    }
22}
23
24struct BencodeSerializer<W: std::io::Write> {
25    writer: W,
26    hack_no_bytestring_prefix: bool,
27}
28
29impl<W: std::io::Write> BencodeSerializer<W> {
30    pub fn new(writer: W) -> Self {
31        Self {
32            writer,
33            hack_no_bytestring_prefix: false,
34        }
35    }
36    fn write_raw(&mut self, buf: &[u8]) -> Result<(), Error> {
37        self.writer.write_all(buf).map_err(Error::Io)
38    }
39    fn write_fmt(&mut self, fmt: core::fmt::Arguments<'_>) -> Result<(), Error> {
40        self.writer.write_fmt(fmt).map_err(Error::Io)
41    }
42    fn write_byte(&mut self, byte: u8) -> Result<(), Error> {
43        self.write_raw(&[byte])
44    }
45    fn write_number<N: std::fmt::Display>(&mut self, number: N) -> Result<(), Error> {
46        self.write_fmt(format_args!("i{number}e"))
47    }
48    fn write_bytes(&mut self, bytes: &[u8]) -> Result<(), Error> {
49        if !self.hack_no_bytestring_prefix {
50            self.write_fmt(format_args!("{}:", bytes.len()))?;
51        }
52        self.write_raw(bytes)
53    }
54}
55
56struct SerializeSeq<'ser, W: std::io::Write> {
57    ser: &'ser mut BencodeSerializer<W>,
58}
59impl<W: std::io::Write> serde::ser::SerializeSeq for SerializeSeq<'_, W> {
60    type Ok = ();
61
62    type Error = Error;
63
64    fn serialize_element<T>(&mut self, value: &T) -> Result<(), Self::Error>
65    where
66        T: ?Sized + serde::Serialize,
67    {
68        value.serialize(&mut *self.ser)
69    }
70
71    fn end(self) -> Result<Self::Ok, Self::Error> {
72        self.ser.write_byte(b'e')
73    }
74}
75
76struct SerializeTuple<'ser, W: std::io::Write> {
77    ser: &'ser mut BencodeSerializer<W>,
78}
79impl<W: std::io::Write> serde::ser::SerializeTuple for SerializeTuple<'_, W> {
80    type Ok = ();
81
82    type Error = Error;
83
84    fn serialize_element<T>(&mut self, value: &T) -> Result<(), Self::Error>
85    where
86        T: ?Sized + serde::Serialize,
87    {
88        value.serialize(&mut *self.ser)
89    }
90
91    fn end(self) -> Result<Self::Ok, Self::Error> {
92        self.ser.write_byte(b'e')
93    }
94}
95
96struct SerializeMap<'ser, W: std::io::Write> {
97    ser: &'ser mut BencodeSerializer<W>,
98    tmp: Vec<(ByteBufOwned, ByteBufOwned)>,
99    last_key: Option<ByteBufOwned>,
100}
101impl<W: std::io::Write> serde::ser::SerializeMap for SerializeMap<'_, W> {
102    type Ok = ();
103
104    type Error = Error;
105
106    fn serialize_key<T>(&mut self, key: &T) -> Result<(), Self::Error>
107    where
108        T: ?Sized + serde::Serialize,
109    {
110        let mut buf = Vec::new();
111        let mut ser = BencodeSerializer::new(&mut buf);
112        ser.hack_no_bytestring_prefix = true;
113        key.serialize(&mut ser)?;
114        self.last_key.replace(ByteBufOwned::from(buf));
115        Ok(())
116    }
117
118    fn serialize_value<T>(&mut self, value: &T) -> Result<(), Self::Error>
119    where
120        T: ?Sized + serde::Serialize,
121    {
122        let mut buf = Vec::new();
123        let mut ser = BencodeSerializer::new(&mut buf);
124        value.serialize(&mut ser)?;
125        self.tmp
126            .push((self.last_key.take().unwrap(), ByteBufOwned::from(buf)));
127        Ok(())
128    }
129
130    fn end(mut self) -> Result<Self::Ok, Self::Error> {
131        self.tmp
132            .sort_unstable_by(|(left, _), (right, _)| left.cmp(right));
133        for (key, value) in self.tmp {
134            self.ser.write_bytes(key.as_ref())?;
135            self.ser.write_raw(value.as_ref())?;
136        }
137        self.ser.write_byte(b'e')
138    }
139}
140
141struct SerializeStruct<'ser, W: std::io::Write> {
142    ser: &'ser mut BencodeSerializer<W>,
143    tmp: Vec<(&'static str, ByteBufOwned)>,
144}
145impl<W: std::io::Write> serde::ser::SerializeStruct for SerializeStruct<'_, W> {
146    type Ok = ();
147
148    type Error = Error;
149
150    fn serialize_field<T>(&mut self, key: &'static str, value: &T) -> Result<(), Self::Error>
151    where
152        T: ?Sized + serde::Serialize,
153    {
154        let mut buf = Vec::new();
155        let mut ser = BencodeSerializer::new(&mut buf);
156        value.serialize(&mut ser)?;
157        self.tmp.push((key, ByteBufOwned::from(buf)));
158        Ok(())
159    }
160
161    fn end(mut self) -> Result<Self::Ok, Self::Error> {
162        self.tmp.sort_unstable_by_key(|&(k, _)| k);
163        for (key, value) in self.tmp {
164            self.ser.write_bytes(key.as_bytes())?;
165            self.ser.write_raw(value.as_ref())?;
166        }
167        self.ser.write_byte(b'e')
168    }
169}
170
171impl<'ser, W: std::io::Write> Serializer for &'ser mut BencodeSerializer<W> {
172    type Ok = ();
173    type Error = Error;
174    type SerializeSeq = SerializeSeq<'ser, W>;
175    type SerializeTuple = SerializeTuple<'ser, W>;
176    type SerializeTupleStruct = Impossible<(), Error>;
177    type SerializeTupleVariant = Impossible<(), Error>;
178    type SerializeMap = SerializeMap<'ser, W>;
179    type SerializeStruct = SerializeStruct<'ser, W>;
180    type SerializeStructVariant = Impossible<(), Error>;
181
182    fn serialize_bool(self, value: bool) -> Result<Self::Ok, Self::Error> {
183        self.write_number(if value { 1 } else { 0 })
184    }
185
186    fn serialize_i8(self, v: i8) -> Result<Self::Ok, Self::Error> {
187        self.write_number(v)
188    }
189
190    fn serialize_i16(self, v: i16) -> Result<Self::Ok, Self::Error> {
191        self.write_number(v)
192    }
193
194    fn serialize_i32(self, v: i32) -> Result<Self::Ok, Self::Error> {
195        self.write_number(v)
196    }
197
198    fn serialize_i64(self, v: i64) -> Result<Self::Ok, Self::Error> {
199        self.write_number(v)
200    }
201
202    fn serialize_u8(self, v: u8) -> Result<Self::Ok, Self::Error> {
203        self.write_number(v)
204    }
205
206    fn serialize_u16(self, v: u16) -> Result<Self::Ok, Self::Error> {
207        self.write_number(v)
208    }
209
210    fn serialize_u32(self, v: u32) -> Result<Self::Ok, Self::Error> {
211        self.write_number(v)
212    }
213
214    fn serialize_u64(self, v: u64) -> Result<Self::Ok, Self::Error> {
215        self.write_number(v)
216    }
217
218    fn serialize_f32(self, _: f32) -> Result<Self::Ok, Self::Error> {
219        Err(Error::Text(&"bencode doesn't support f32"))
220    }
221
222    fn serialize_f64(self, _: f64) -> Result<Self::Ok, Self::Error> {
223        Err(Error::Text(&"bencode doesn't support f32"))
224    }
225
226    fn serialize_char(self, _: char) -> Result<Self::Ok, Self::Error> {
227        Err(Error::Text(&"bencode doesn't support chars"))
228    }
229
230    fn serialize_str(self, v: &str) -> Result<Self::Ok, Self::Error> {
231        self.write_bytes(v.as_bytes())
232    }
233
234    fn serialize_bytes(self, v: &[u8]) -> Result<Self::Ok, Self::Error> {
235        self.write_bytes(v)
236    }
237
238    fn serialize_none(self) -> Result<Self::Ok, Self::Error> {
239        Err(Error::Text(&"bencode doesn't support None"))
240    }
241
242    fn serialize_some<T>(self, value: &T) -> Result<Self::Ok, Self::Error>
243    where
244        T: ?Sized + serde::Serialize,
245    {
246        value.serialize(self)
247    }
248
249    fn serialize_unit(self) -> Result<Self::Ok, Self::Error> {
250        Err(Error::Text(&"bencode doesn't support Rust unit ()"))
251    }
252
253    fn serialize_unit_struct(self, _name: &'static str) -> Result<Self::Ok, Self::Error> {
254        Err(Error::Text(&"bencode doesn't support unit structs"))
255    }
256
257    fn serialize_unit_variant(
258        self,
259        _name: &'static str,
260        _variant_index: u32,
261        _variant: &'static str,
262    ) -> Result<Self::Ok, Self::Error> {
263        Err(Error::Text(&"bencode doesn't support unit variants"))
264    }
265
266    fn serialize_newtype_struct<T>(
267        self,
268        name: &'static str,
269        value: &T,
270    ) -> Result<Self::Ok, Self::Error>
271    where
272        T: ?Sized + serde::Serialize,
273    {
274        if name == crate::raw_value::TAG {
275            self.hack_no_bytestring_prefix = true;
276            value.serialize(&mut *self)?;
277            self.hack_no_bytestring_prefix = false;
278            return Ok(());
279        }
280        Err(Error::Text(&"bencode doesn't support newtype structs"))
281    }
282
283    fn serialize_newtype_variant<T>(
284        self,
285        _name: &'static str,
286        _variant_index: u32,
287        _variant: &'static str,
288        _value: &T,
289    ) -> Result<Self::Ok, Self::Error>
290    where
291        T: ?Sized + serde::Serialize,
292    {
293        Err(Error::Text(&"bencode doesn't support newtype variants"))
294    }
295
296    fn serialize_seq(self, _len: Option<usize>) -> Result<Self::SerializeSeq, Self::Error> {
297        self.write_byte(b'l')?;
298        Ok(SerializeSeq { ser: self })
299    }
300
301    fn serialize_tuple(self, _len: usize) -> Result<Self::SerializeTuple, Self::Error> {
302        Err(Error::Text(&"bencode doesn't support tuples"))
303    }
304
305    fn serialize_tuple_struct(
306        self,
307        _name: &'static str,
308        _len: usize,
309    ) -> Result<Self::SerializeTupleStruct, Self::Error> {
310        Err(Error::Text(&"bencode doesn't support tuple structs"))
311    }
312
313    fn serialize_tuple_variant(
314        self,
315        _name: &'static str,
316        _variant_index: u32,
317        _variant: &'static str,
318        _len: usize,
319    ) -> Result<Self::SerializeTupleVariant, Self::Error> {
320        Err(Error::Text(&"bencode doesn't support tuple variants"))
321    }
322
323    fn serialize_map(self, _len: Option<usize>) -> Result<Self::SerializeMap, Self::Error> {
324        self.write_byte(b'd')?;
325        Ok(SerializeMap {
326            ser: self,
327            tmp: Default::default(),
328            last_key: None,
329        })
330    }
331
332    fn serialize_struct(
333        self,
334        _name: &'static str,
335        len: usize,
336    ) -> Result<Self::SerializeStruct, Self::Error> {
337        self.write_byte(b'd')?;
338        Ok(SerializeStruct {
339            ser: self,
340            tmp: Vec::with_capacity(len),
341        })
342    }
343
344    fn serialize_struct_variant(
345        self,
346        _name: &'static str,
347        _variant_index: u32,
348        _variant: &'static str,
349        _len: usize,
350    ) -> Result<Self::SerializeStructVariant, Self::Error> {
351        Err(Error::Text(&"bencode doesn't support struct variants"))
352    }
353}
354
355pub fn bencode_serialize_to_writer<T: Serialize, W: std::io::Write>(
356    value: T,
357    writer: &mut W,
358) -> Result<(), Error> {
359    let mut serializer = BencodeSerializer::new(writer);
360    value.serialize(&mut serializer)?;
361    Ok(())
362}
363
364#[cfg(test)]
365mod tests {
366    use buffers::{ByteBuf, ByteBufOwned};
367    use serde::Serialize;
368
369    use crate::bencode_serialize_to_writer;
370
371    fn ser<T: Serialize>(value: T) -> Result<ByteBufOwned, crate::SerializeError> {
372        let mut vec = Vec::new();
373        bencode_serialize_to_writer(&value, &mut vec)?;
374        Ok(vec.into())
375    }
376
377    #[test]
378    fn test_ints() {
379        assert_eq!(ser(42u16).unwrap(), b"i42e"[..].into());
380        assert_eq!(ser(42u32).unwrap(), b"i42e"[..].into());
381        assert_eq!(ser(42u64).unwrap(), b"i42e"[..].into());
382    }
383
384    #[test]
385    fn test_bytes() {
386        assert_eq!(ser(ByteBuf(b"abc")).unwrap(), b"3:abc"[..].into());
387        assert_eq!(
388            ser(ByteBufOwned::from(&b"abc"[..])).unwrap(),
389            b"3:abc"[..].into()
390        );
391    }
392
393    #[test]
394    fn test_seq() {
395        assert_eq!(
396            ser(&[ByteBuf(b"foo"), ByteBuf(b"bar")][..]).unwrap(),
397            b"l3:foo3:bare"[..].into()
398        );
399        assert_eq!(
400            ser(vec![ByteBuf(b"foo"), ByteBuf(b"bar")]).unwrap(),
401            b"l3:foo3:bare"[..].into()
402        );
403    }
404
405    #[test]
406    fn test_struct() {
407        #[derive(serde_derive::Serialize, Debug)]
408        struct S<'a> {
409            key: u32,
410            value: ByteBuf<'a>,
411        }
412        assert_eq!(
413            ser(S {
414                key: 42,
415                value: b"foo"[..].into()
416            })
417            .unwrap(),
418            b"d3:keyi42e5:value3:fooe"[..].into()
419        );
420    }
421
422    #[test]
423    fn test_dict() {
424        use std::collections::BTreeMap;
425        let mut m = BTreeMap::new();
426        m.insert(ByteBuf(b"key"), ByteBuf(b"value"));
427        m.insert(ByteBuf(b"key2"), ByteBuf(b"value2"));
428        assert_eq!(ser(&m).unwrap(), b"d3:key5:value4:key26:value2e"[..].into());
429    }
430
431    #[test]
432    fn test_struct_with_option() {
433        #[derive(serde_derive::Serialize, Debug)]
434        struct S<'a> {
435            key: u32,
436            #[serde(skip_serializing_if = "Option::is_none")]
437            value: Option<ByteBuf<'a>>,
438        }
439        assert_eq!(
440            ser(S {
441                key: 42,
442                value: Some(b"foo"[..].into())
443            })
444            .unwrap(),
445            b"d3:keyi42e5:value3:fooe"[..].into()
446        );
447        assert_eq!(
448            ser(S {
449                key: 42,
450                value: None
451            })
452            .unwrap(),
453            b"d3:keyi42ee"[..].into()
454        );
455    }
456}