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