Skip to main content

irontide_bencode/
ser.rs

1use serde::ser::{self, Serialize};
2use std::io::Write;
3
4use crate::error::{Error, Result};
5
6/// Bencode serializer.
7///
8/// Writes bencode-encoded data to any `std::io::Write` sink.
9pub struct Serializer<W> {
10    writer: W,
11}
12
13impl<W: Write> Serializer<W> {
14    /// Create a new serializer writing to the given sink.
15    pub fn new(writer: W) -> Self {
16        Serializer { writer }
17    }
18
19    fn write_byte_string(&mut self, v: &[u8]) -> Result<()> {
20        write!(self.writer, "{}:", v.len()).map_err(|e| Error::Custom(e.to_string()))?;
21        self.writer
22            .write_all(v)
23            .map_err(|e| Error::Custom(e.to_string()))
24    }
25
26    fn write_integer(&mut self, v: i64) -> Result<()> {
27        write!(self.writer, "i{v}e").map_err(|e| Error::Custom(e.to_string()))
28    }
29}
30
31impl<'a, W: Write> ser::Serializer for &'a mut Serializer<W> {
32    type Ok = ();
33    type Error = Error;
34    type SerializeSeq = &'a mut Serializer<W>;
35    type SerializeTuple = &'a mut Serializer<W>;
36    type SerializeTupleStruct = &'a mut Serializer<W>;
37    type SerializeTupleVariant = &'a mut Serializer<W>;
38    type SerializeMap = SortedMapSerializer<'a, W>;
39    type SerializeStruct = SortedMapSerializer<'a, W>;
40    type SerializeStructVariant = SortedMapSerializer<'a, W>;
41
42    fn serialize_bool(self, v: bool) -> Result<()> {
43        self.write_integer(if v { 1 } else { 0 })
44    }
45
46    fn serialize_i8(self, v: i8) -> Result<()> {
47        self.write_integer(v as i64)
48    }
49
50    fn serialize_i16(self, v: i16) -> Result<()> {
51        self.write_integer(v as i64)
52    }
53
54    fn serialize_i32(self, v: i32) -> Result<()> {
55        self.write_integer(v as i64)
56    }
57
58    fn serialize_i64(self, v: i64) -> Result<()> {
59        self.write_integer(v)
60    }
61
62    fn serialize_u8(self, v: u8) -> Result<()> {
63        self.write_integer(v as i64)
64    }
65
66    fn serialize_u16(self, v: u16) -> Result<()> {
67        self.write_integer(v as i64)
68    }
69
70    fn serialize_u32(self, v: u32) -> Result<()> {
71        self.write_integer(v as i64)
72    }
73
74    fn serialize_u64(self, v: u64) -> Result<()> {
75        self.write_integer(v as i64)
76    }
77
78    fn serialize_f32(self, _v: f32) -> Result<()> {
79        Err(Error::Custom("bencode does not support floats".into()))
80    }
81
82    fn serialize_f64(self, _v: f64) -> Result<()> {
83        Err(Error::Custom("bencode does not support floats".into()))
84    }
85
86    fn serialize_char(self, v: char) -> Result<()> {
87        let mut buf = [0u8; 4];
88        let s = v.encode_utf8(&mut buf);
89        self.write_byte_string(s.as_bytes())
90    }
91
92    fn serialize_str(self, v: &str) -> Result<()> {
93        self.write_byte_string(v.as_bytes())
94    }
95
96    fn serialize_bytes(self, v: &[u8]) -> Result<()> {
97        self.write_byte_string(v)
98    }
99
100    fn serialize_none(self) -> Result<()> {
101        Err(Error::Custom("bencode does not support None/null".into()))
102    }
103
104    fn serialize_some<T: ?Sized + Serialize>(self, value: &T) -> Result<()> {
105        value.serialize(self)
106    }
107
108    fn serialize_unit(self) -> Result<()> {
109        Err(Error::Custom("bencode does not support unit".into()))
110    }
111
112    fn serialize_unit_struct(self, _name: &'static str) -> Result<()> {
113        Err(Error::Custom(
114            "bencode does not support unit structs".into(),
115        ))
116    }
117
118    fn serialize_unit_variant(
119        self,
120        _name: &'static str,
121        _variant_index: u32,
122        variant: &'static str,
123    ) -> Result<()> {
124        self.write_byte_string(variant.as_bytes())
125    }
126
127    fn serialize_newtype_struct<T: ?Sized + Serialize>(
128        self,
129        _name: &'static str,
130        value: &T,
131    ) -> Result<()> {
132        value.serialize(self)
133    }
134
135    fn serialize_newtype_variant<T: ?Sized + Serialize>(
136        self,
137        _name: &'static str,
138        _variant_index: u32,
139        variant: &'static str,
140        value: &T,
141    ) -> Result<()> {
142        self.writer
143            .write_all(b"d")
144            .map_err(|e| Error::Custom(e.to_string()))?;
145        self.write_byte_string(variant.as_bytes())?;
146        value.serialize(&mut *self)?;
147        self.writer
148            .write_all(b"e")
149            .map_err(|e| Error::Custom(e.to_string()))
150    }
151
152    fn serialize_seq(self, _len: Option<usize>) -> Result<Self::SerializeSeq> {
153        self.writer
154            .write_all(b"l")
155            .map_err(|e| Error::Custom(e.to_string()))?;
156        Ok(self)
157    }
158
159    fn serialize_tuple(self, len: usize) -> Result<Self::SerializeTuple> {
160        self.serialize_seq(Some(len))
161    }
162
163    fn serialize_tuple_struct(
164        self,
165        _name: &'static str,
166        len: usize,
167    ) -> Result<Self::SerializeTupleStruct> {
168        self.serialize_seq(Some(len))
169    }
170
171    fn serialize_tuple_variant(
172        self,
173        _name: &'static str,
174        _variant_index: u32,
175        variant: &'static str,
176        _len: usize,
177    ) -> Result<Self::SerializeTupleVariant> {
178        self.writer
179            .write_all(b"d")
180            .map_err(|e| Error::Custom(e.to_string()))?;
181        self.write_byte_string(variant.as_bytes())?;
182        self.writer
183            .write_all(b"l")
184            .map_err(|e| Error::Custom(e.to_string()))?;
185        Ok(self)
186    }
187
188    fn serialize_map(self, _len: Option<usize>) -> Result<Self::SerializeMap> {
189        Ok(SortedMapSerializer::new(self))
190    }
191
192    fn serialize_struct(self, _name: &'static str, _len: usize) -> Result<Self::SerializeStruct> {
193        Ok(SortedMapSerializer::new(self))
194    }
195
196    fn serialize_struct_variant(
197        self,
198        _name: &'static str,
199        _variant_index: u32,
200        variant: &'static str,
201        _len: usize,
202    ) -> Result<Self::SerializeStructVariant> {
203        // Wrap in a dict: d<variant><dict-value>e
204        // The outer dict is handled by the caller; we buffer this inner dict.
205        // Actually, for struct variants we need to write the outer dict key.
206        self.writer
207            .write_all(b"d")
208            .map_err(|e| Error::Custom(e.to_string()))?;
209        self.write_byte_string(variant.as_bytes())?;
210        Ok(SortedMapSerializer::new_nested(self))
211    }
212}
213
214// Sequence serialization
215impl<W: Write> ser::SerializeSeq for &mut Serializer<W> {
216    type Ok = ();
217    type Error = Error;
218
219    fn serialize_element<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<()> {
220        value.serialize(&mut **self)
221    }
222
223    fn end(self) -> Result<()> {
224        self.writer
225            .write_all(b"e")
226            .map_err(|e| Error::Custom(e.to_string()))
227    }
228}
229
230impl<W: Write> ser::SerializeTuple for &mut Serializer<W> {
231    type Ok = ();
232    type Error = Error;
233
234    fn serialize_element<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<()> {
235        ser::SerializeSeq::serialize_element(self, value)
236    }
237
238    fn end(self) -> Result<()> {
239        ser::SerializeSeq::end(self)
240    }
241}
242
243impl<W: Write> ser::SerializeTupleStruct for &mut Serializer<W> {
244    type Ok = ();
245    type Error = Error;
246
247    fn serialize_field<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<()> {
248        ser::SerializeSeq::serialize_element(self, value)
249    }
250
251    fn end(self) -> Result<()> {
252        ser::SerializeSeq::end(self)
253    }
254}
255
256impl<W: Write> ser::SerializeTupleVariant for &mut Serializer<W> {
257    type Ok = ();
258    type Error = Error;
259
260    fn serialize_field<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<()> {
261        ser::SerializeSeq::serialize_element(self, value)
262    }
263
264    fn end(self) -> Result<()> {
265        // Close list, then close outer dict
266        self.writer
267            .write_all(b"e")
268            .map_err(|e| Error::Custom(e.to_string()))?;
269        self.writer
270            .write_all(b"e")
271            .map_err(|e| Error::Custom(e.to_string()))
272    }
273}
274
275/// Buffering map serializer that sorts keys before writing.
276///
277/// BEP 3 requires dictionary keys to be in sorted order. We buffer all
278/// key-value pairs, sort by key, then write them all at once in `end()`.
279pub struct SortedMapSerializer<'a, W> {
280    ser: &'a mut Serializer<W>,
281    entries: Vec<(Vec<u8>, Vec<u8>)>,
282    current_key: Option<Vec<u8>>,
283    /// If true, we need to close an outer dict wrapper on end().
284    nested: bool,
285}
286
287impl<'a, W: Write> SortedMapSerializer<'a, W> {
288    fn new(ser: &'a mut Serializer<W>) -> Self {
289        SortedMapSerializer {
290            ser,
291            entries: Vec::new(),
292            current_key: None,
293            nested: false,
294        }
295    }
296
297    fn new_nested(ser: &'a mut Serializer<W>) -> Self {
298        SortedMapSerializer {
299            ser,
300            entries: Vec::new(),
301            current_key: None,
302            nested: true,
303        }
304    }
305
306    fn flush_sorted(&mut self) -> Result<()> {
307        self.entries.sort_by(|a, b| a.0.cmp(&b.0));
308        self.ser
309            .writer
310            .write_all(b"d")
311            .map_err(|e| Error::Custom(e.to_string()))?;
312        for (key, value) in &self.entries {
313            self.ser.write_byte_string(key)?;
314            self.ser
315                .writer
316                .write_all(value)
317                .map_err(|e| Error::Custom(e.to_string()))?;
318        }
319        self.ser
320            .writer
321            .write_all(b"e")
322            .map_err(|e| Error::Custom(e.to_string()))
323    }
324}
325
326impl<W: Write> ser::SerializeMap for SortedMapSerializer<'_, W> {
327    type Ok = ();
328    type Error = Error;
329
330    fn serialize_key<T: ?Sized + Serialize>(&mut self, key: &T) -> Result<()> {
331        let mut buf = Vec::new();
332        let mut key_ser = Serializer::new(&mut buf);
333        key.serialize(&mut key_ser)?;
334        // Extract the raw key bytes from the serialized bencode string.
335        // Keys serialize as `<len>:<bytes>`, so we need the raw bytes.
336        self.current_key = Some(extract_string_content(&buf)?);
337        Ok(())
338    }
339
340    fn serialize_value<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<()> {
341        let key = self
342            .current_key
343            .take()
344            .ok_or_else(|| Error::Custom("serialize_value called before serialize_key".into()))?;
345        let mut buf = Vec::new();
346        let mut val_ser = Serializer::new(&mut buf);
347        value.serialize(&mut val_ser)?;
348        self.entries.push((key, buf));
349        Ok(())
350    }
351
352    fn end(mut self) -> Result<()> {
353        self.flush_sorted()?;
354        if self.nested {
355            self.ser
356                .writer
357                .write_all(b"e")
358                .map_err(|e| Error::Custom(e.to_string()))?;
359        }
360        Ok(())
361    }
362}
363
364impl<W: Write> ser::SerializeStruct for SortedMapSerializer<'_, W> {
365    type Ok = ();
366    type Error = Error;
367
368    fn serialize_field<T: ?Sized + Serialize>(
369        &mut self,
370        key: &'static str,
371        value: &T,
372    ) -> Result<()> {
373        let mut buf = Vec::new();
374        let mut val_ser = Serializer::new(&mut buf);
375        value.serialize(&mut val_ser)?;
376        self.entries.push((key.as_bytes().to_vec(), buf));
377        Ok(())
378    }
379
380    fn end(mut self) -> Result<()> {
381        self.flush_sorted()?;
382        if self.nested {
383            self.ser
384                .writer
385                .write_all(b"e")
386                .map_err(|e| Error::Custom(e.to_string()))?;
387        }
388        Ok(())
389    }
390
391    fn skip_field(&mut self, _key: &'static str) -> Result<()> {
392        Ok(())
393    }
394}
395
396impl<W: Write> ser::SerializeStructVariant for SortedMapSerializer<'_, W> {
397    type Ok = ();
398    type Error = Error;
399
400    fn serialize_field<T: ?Sized + Serialize>(
401        &mut self,
402        key: &'static str,
403        value: &T,
404    ) -> Result<()> {
405        ser::SerializeStruct::serialize_field(self, key, value)
406    }
407
408    fn end(self) -> Result<()> {
409        ser::SerializeStruct::end(self)
410    }
411}
412
413/// Extract the raw byte content from a bencode-encoded string `<len>:<data>`.
414fn extract_string_content(encoded: &[u8]) -> Result<Vec<u8>> {
415    let colon_pos = encoded
416        .iter()
417        .position(|&b| b == b':')
418        .ok_or_else(|| Error::Custom("map key must be a string".into()))?;
419    Ok(encoded[colon_pos + 1..].to_vec())
420}
421
422#[cfg(test)]
423mod tests {
424    use serde::Serialize;
425
426    #[test]
427    fn serialize_integer() {
428        assert_eq!(to_bytes(&42i64), b"i42e");
429        assert_eq!(to_bytes(&0i64), b"i0e");
430        assert_eq!(to_bytes(&-1i64), b"i-1e");
431    }
432
433    #[test]
434    fn serialize_string() {
435        assert_eq!(to_bytes("spam"), b"4:spam");
436        assert_eq!(to_bytes(""), b"0:");
437    }
438
439    #[test]
440    fn serialize_list() {
441        let list = vec!["spam", "eggs"];
442        assert_eq!(to_bytes(&list), b"l4:spam4:eggse");
443    }
444
445    #[test]
446    fn serialize_struct_sorted_keys() {
447        #[derive(Serialize)]
448        struct Test {
449            zebra: i64,
450            alpha: i64,
451        }
452        let val = Test { zebra: 2, alpha: 1 };
453        // Keys must be sorted: alpha before zebra
454        assert_eq!(to_bytes(&val), b"d5:alphai1e5:zebrai2ee");
455    }
456
457    fn to_bytes<T: Serialize + ?Sized>(value: &T) -> Vec<u8> {
458        crate::to_bytes(value).unwrap()
459    }
460}