Skip to main content

librqbit_bencode/
bencode_value.rs

1use std::{collections::HashMap, fmt::Display, marker::PhantomData};
2
3use buffers::{ByteBuf, ByteBufOwned};
4use bytes::Bytes;
5use clone_to_owned::CloneToOwned;
6use serde::{Deserialize, Deserializer, Serialize};
7
8use crate::deserialize::from_bytes;
9
10pub fn dyn_from_bytes<'de, BufT>(
11    buf: &'de [u8],
12) -> Result<BencodeValue<BufT>, crate::deserialize::ErrorWithContext<'de>>
13where
14    BufT: From<&'de [u8]> + std::hash::Hash + Eq,
15{
16    from_bytes(buf)
17}
18
19impl<BufT: serde::Serialize + Eq + std::hash::Hash> serde::Serialize for BencodeValue<BufT> {
20    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
21    where
22        S: serde::Serializer,
23    {
24        match self {
25            BencodeValue::Bytes(b) => b.serialize(serializer),
26            BencodeValue::Integer(v) => v.serialize(serializer),
27            BencodeValue::List(l) => l.serialize(serializer),
28            BencodeValue::Dict(d) => d.serialize(serializer),
29        }
30    }
31}
32
33impl<'de, BufT> serde::de::Deserialize<'de> for BencodeValue<BufT>
34where
35    BufT: From<&'de [u8]> + std::hash::Hash + Eq,
36{
37    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
38    where
39        D: Deserializer<'de>,
40    {
41        struct Visitor<BufT> {
42            buftype: PhantomData<BufT>,
43        }
44
45        impl<'de, BufT> serde::de::Visitor<'de> for Visitor<BufT>
46        where
47            BufT: From<&'de [u8]> + std::hash::Hash + Eq,
48        {
49            type Value = BencodeValue<BufT>;
50
51            fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
52                write!(formatter, "a bencode value")
53            }
54
55            fn visit_i64<E>(self, v: i64) -> Result<Self::Value, E>
56            where
57                E: serde::de::Error,
58            {
59                Ok(BencodeValue::Integer(v))
60            }
61
62            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
63            where
64                A: serde::de::SeqAccess<'de>,
65            {
66                let mut v = Vec::new();
67                while let Some(value) = seq.next_element()? {
68                    v.push(value);
69                }
70                Ok(BencodeValue::List(v))
71            }
72
73            fn visit_borrowed_bytes<E>(self, v: &'de [u8]) -> Result<Self::Value, E>
74            where
75                E: serde::de::Error,
76            {
77                Ok(BencodeValue::Bytes(BufT::from(v)))
78            }
79
80            fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
81            where
82                A: serde::de::MapAccess<'de>,
83            {
84                let mut hmap = HashMap::new();
85                while let Some(key) = map.next_key::<&'de [u8]>()? {
86                    let value = map.next_value()?;
87                    hmap.insert(BufT::from(key), value);
88                }
89                Ok(BencodeValue::Dict(hmap))
90            }
91        }
92
93        deserializer.deserialize_any(Visitor {
94            buftype: PhantomData,
95        })
96    }
97}
98
99// A dynamic value when we don't know exactly what we are deserializing.
100// Useful for debugging.
101
102#[derive(PartialEq, Eq)]
103pub enum BencodeValue<BufT: std::hash::Hash + Eq> {
104    Bytes(BufT),
105    Integer(i64),
106    List(Vec<BencodeValue<BufT>>),
107    Dict(HashMap<BufT, BencodeValue<BufT>>),
108}
109
110impl<BufT: std::fmt::Debug + std::hash::Hash + Eq> std::fmt::Debug for BencodeValue<BufT> {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        match self {
113            BencodeValue::Bytes(b) => std::fmt::Debug::fmt(b, f),
114            BencodeValue::Integer(i) => std::fmt::Debug::fmt(i, f),
115            BencodeValue::List(l) => std::fmt::Debug::fmt(l, f),
116            BencodeValue::Dict(d) => std::fmt::Debug::fmt(d, f),
117        }
118    }
119}
120
121impl<BufT> CloneToOwned for BencodeValue<BufT>
122where
123    BufT: CloneToOwned + std::hash::Hash + Eq,
124    <BufT as CloneToOwned>::Target: Eq + std::hash::Hash,
125{
126    type Target = BencodeValue<<BufT as CloneToOwned>::Target>;
127
128    fn clone_to_owned(&self, within_buffer: Option<&Bytes>) -> Self::Target {
129        match self {
130            BencodeValue::Bytes(b) => BencodeValue::Bytes(b.clone_to_owned(within_buffer)),
131            BencodeValue::Integer(i) => BencodeValue::Integer(*i),
132            BencodeValue::List(l) => BencodeValue::List(l.clone_to_owned(within_buffer)),
133            BencodeValue::Dict(d) => BencodeValue::Dict(d.clone_to_owned(within_buffer)),
134        }
135    }
136}
137
138pub type BencodeValueBorrowed<'a> = BencodeValue<ByteBuf<'a>>;
139pub type BencodeValueOwned = BencodeValue<ByteBufOwned>;
140
141// A wrapper to deserialize dyn values as strings.
142#[derive(PartialEq, Eq, Hash)]
143pub struct AsDisplay<T>(T);
144
145impl<'de, T> From<&'de [u8]> for AsDisplay<T>
146where
147    T: From<&'de [u8]>,
148{
149    fn from(value: &'de [u8]) -> Self {
150        Self(T::from(value))
151    }
152}
153
154impl<T> Serialize for AsDisplay<T>
155where
156    T: Display,
157{
158    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
159    where
160        S: serde::Serializer,
161    {
162        serializer.serialize_str(&format!("{}", self.0))
163    }
164}
165
166impl<'de, T> Deserialize<'de> for AsDisplay<T>
167where
168    T: Deserialize<'de>,
169{
170    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
171    where
172        D: Deserializer<'de>,
173    {
174        let value = T::deserialize(deserializer)?;
175        Ok(AsDisplay(value))
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    use crate::serialize::bencode_serialize_to_writer;
182
183    use super::*;
184    use serde::Serialize;
185    use std::io::Read;
186
187    #[test]
188    fn test_deserialize_torrent_dyn() {
189        let mut buf = Vec::new();
190        let filename = "../librqbit/resources/ubuntu-21.04-desktop-amd64.iso.torrent";
191        std::fs::File::open(filename)
192            .unwrap()
193            .read_to_end(&mut buf)
194            .unwrap();
195
196        let torrent_borrowed: BencodeValueBorrowed = from_bytes(&buf).unwrap();
197        let torrent_owned: BencodeValueOwned = from_bytes(&buf).unwrap();
198        dbg!(torrent_borrowed);
199        dbg!(torrent_owned);
200    }
201
202    #[test]
203    fn test_serialize_torrent_dyn() {
204        let mut buf = Vec::new();
205        let filename = "../librqbit/resources/ubuntu-21.04-desktop-amd64.iso.torrent";
206        std::fs::File::open(filename)
207            .unwrap()
208            .read_to_end(&mut buf)
209            .unwrap();
210
211        let torrent: BencodeValueBorrowed = from_bytes(&buf).unwrap();
212
213        let mut buf = Vec::<u8>::new();
214        bencode_serialize_to_writer(&torrent, &mut buf).unwrap();
215
216        let new_torrent = from_bytes(&buf).unwrap();
217        assert_eq!(torrent, new_torrent);
218    }
219
220    #[test]
221    fn test_serialize_struct_with_option() {
222        #[derive(Serialize)]
223        struct Test {
224            f1: i64,
225            #[serde(skip_serializing_if = "Option::is_none")]
226            missing: Option<i64>,
227        }
228        let test = Test {
229            f1: 100,
230            missing: None,
231        };
232        let mut buf = Vec::<u8>::new();
233        bencode_serialize_to_writer(&test, &mut buf).unwrap();
234        assert_eq!(&buf, b"d2:f1i100ee");
235    }
236}