dg_xch_serialize/
lib.rs

1use log::warn;
2use std::collections::HashMap;
3use std::convert::Infallible;
4use std::fmt::{Display, Formatter};
5use std::hash::Hash;
6use std::io::{Cursor, Error, ErrorKind, Read, Write};
7use std::str::FromStr;
8use time::OffsetDateTime;
9
10#[derive(
11    Default,
12    Debug,
13    Copy,
14    Clone,
15    Ord,
16    PartialOrd,
17    Eq,
18    PartialEq,
19    serde::Serialize,
20    serde::Deserialize,
21)]
22pub enum ChiaProtocolVersion {
23    Chia0_0_34 = 34, //Pre 2.0.0
24    Chia0_0_35 = 35, //2.0.0
25    Chia0_0_36 = 36, //2.2.0
26    #[default]
27    Chia0_0_37 = 37, //2.5.5
28}
29impl Display for ChiaProtocolVersion {
30    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
31        match self {
32            ChiaProtocolVersion::Chia0_0_34 => f.write_str("0.0.34"),
33            ChiaProtocolVersion::Chia0_0_35 => f.write_str("0.0.35"),
34            ChiaProtocolVersion::Chia0_0_36 => f.write_str("0.0.36"),
35            ChiaProtocolVersion::Chia0_0_37 => f.write_str("0.0.37"),
36        }
37    }
38}
39impl FromStr for ChiaProtocolVersion {
40    type Err = Infallible;
41    fn from_str(s: &str) -> Result<Self, Self::Err> {
42        Ok(match s {
43            "0.0.34" => ChiaProtocolVersion::Chia0_0_34,
44            "0.0.35" => ChiaProtocolVersion::Chia0_0_35,
45            "0.0.36" => ChiaProtocolVersion::Chia0_0_36,
46            "0.0.37" => ChiaProtocolVersion::Chia0_0_37,
47            _ => {
48                warn!(
49                    "Failed to detect Protocol Version: {s}, defaulting to {}",
50                    ChiaProtocolVersion::default()
51                );
52                ChiaProtocolVersion::default()
53            }
54        })
55    }
56}
57
58pub trait ChiaSerialize {
59    fn to_bytes(&self, version: ChiaProtocolVersion) -> Result<Vec<u8>, Error>
60    where
61        Self: Sized;
62    fn from_bytes<T: AsRef<[u8]>>(
63        bytes: &mut Cursor<T>,
64        version: ChiaProtocolVersion,
65    ) -> Result<Self, Error>
66    where
67        Self: Sized;
68}
69impl ChiaSerialize for OffsetDateTime {
70    fn to_bytes(&self, version: ChiaProtocolVersion) -> Result<Vec<u8>, Error>
71    where
72        Self: Sized,
73    {
74        (self.unix_timestamp() as u64).to_bytes(version)
75    }
76    fn from_bytes<T: AsRef<[u8]>>(
77        bytes: &mut Cursor<T>,
78        version: ChiaProtocolVersion,
79    ) -> Result<Self, Error>
80    where
81        Self: Sized,
82    {
83        let timestamp: u64 = u64::from_bytes(bytes, version)?;
84        OffsetDateTime::from_unix_timestamp(timestamp as i64)
85            .map_err(|e| Error::new(ErrorKind::InvalidData, e))
86    }
87}
88
89impl ChiaSerialize for String {
90    fn to_bytes(&self, _version: ChiaProtocolVersion) -> Result<Vec<u8>, Error>
91    where
92        Self: Sized,
93    {
94        let mut bytes: Vec<u8> = Vec::new();
95        #[allow(clippy::cast_possible_truncation)]
96        bytes.extend((self.len() as u32).to_be_bytes());
97        bytes.extend(self.as_bytes());
98        Ok(bytes)
99    }
100    fn from_bytes<T: AsRef<[u8]>>(
101        bytes: &mut Cursor<T>,
102        _version: ChiaProtocolVersion,
103    ) -> Result<Self, Error>
104    where
105        Self: Sized,
106    {
107        let mut u32_len_ary: [u8; 4] = [0; 4];
108        bytes.read_exact(&mut u32_len_ary)?;
109        let vec_len = u32::from_be_bytes(u32_len_ary) as usize;
110        if vec_len > 2048 {
111            warn!("Serializing Large Vec: {vec_len}");
112        }
113        let mut buf = vec![0u8; vec_len];
114        bytes.read_exact(&mut buf[0..vec_len])?;
115        String::from_utf8(buf).map_err(|e| {
116            Error::new(
117                ErrorKind::InvalidInput,
118                format!("Failed to parse Utf-8 String from Bytes: {e:?}"),
119            )
120        })
121    }
122}
123
124impl ChiaSerialize for bool {
125    fn to_bytes(&self, _version: ChiaProtocolVersion) -> Result<Vec<u8>, Error>
126    where
127        Self: Sized,
128    {
129        Ok(vec![u8::from(*self)])
130    }
131    fn from_bytes<T: AsRef<[u8]>>(
132        bytes: &mut Cursor<T>,
133        _version: ChiaProtocolVersion,
134    ) -> Result<Self, Error>
135    where
136        Self: Sized,
137    {
138        let mut bool_buf: [u8; 1] = [0; 1];
139        bytes.read_exact(&mut bool_buf)?;
140        match bool_buf[0] {
141            0 => Ok(false),
142            1 => Ok(true),
143            _ => Err(Error::new(
144                ErrorKind::InvalidInput,
145                format!("Failed to parse bool, invalid value: {:?}", bool_buf[0]),
146            )),
147        }
148    }
149}
150
151impl<T> ChiaSerialize for Option<T>
152where
153    T: ChiaSerialize,
154{
155    fn to_bytes(&self, version: ChiaProtocolVersion) -> Result<Vec<u8>, Error>
156    where
157        Self: Sized,
158    {
159        let mut bytes: Vec<u8> = Vec::new();
160        match &self {
161            Some(t) => {
162                bytes.push(1u8);
163                bytes.extend(t.to_bytes(version)?);
164            }
165            None => {
166                bytes.push(0u8);
167            }
168        }
169        Ok(bytes)
170    }
171    fn from_bytes<B: AsRef<[u8]>>(
172        bytes: &mut Cursor<B>,
173        version: ChiaProtocolVersion,
174    ) -> Result<Self, Error>
175    where
176        Self: Sized,
177    {
178        let mut bool_buf: [u8; 1] = [0; 1];
179        bytes.read_exact(&mut bool_buf)?;
180        if bool_buf[0] > 0 {
181            Ok(Some(T::from_bytes(bytes, version)?))
182        } else {
183            Ok(None)
184        }
185    }
186}
187
188impl<T, U> ChiaSerialize for (T, U)
189where
190    T: ChiaSerialize,
191    U: ChiaSerialize,
192{
193    fn to_bytes(&self, version: ChiaProtocolVersion) -> Result<Vec<u8>, Error>
194    where
195        Self: Sized,
196    {
197        let mut bytes: Vec<u8> = Vec::new();
198        bytes.extend(self.0.to_bytes(version)?);
199        bytes.extend(self.1.to_bytes(version)?);
200        Ok(bytes)
201    }
202    fn from_bytes<B: AsRef<[u8]>>(
203        bytes: &mut Cursor<B>,
204        version: ChiaProtocolVersion,
205    ) -> Result<Self, Error>
206    where
207        Self: Sized,
208    {
209        let t = T::from_bytes(bytes, version)?;
210        let u = U::from_bytes(bytes, version)?;
211        Ok((t, u))
212    }
213}
214
215impl<T, U, V> ChiaSerialize for (T, U, V)
216where
217    T: ChiaSerialize,
218    U: ChiaSerialize,
219    V: ChiaSerialize,
220{
221    fn to_bytes(&self, version: ChiaProtocolVersion) -> Result<Vec<u8>, Error>
222    where
223        Self: Sized,
224    {
225        let mut bytes: Vec<u8> = Vec::new();
226        bytes.extend(self.0.to_bytes(version)?);
227        bytes.extend(self.1.to_bytes(version)?);
228        bytes.extend(self.2.to_bytes(version)?);
229        Ok(bytes)
230    }
231    fn from_bytes<B: AsRef<[u8]>>(
232        bytes: &mut Cursor<B>,
233        version: ChiaProtocolVersion,
234    ) -> Result<Self, Error>
235    where
236        Self: Sized,
237    {
238        let t = T::from_bytes(bytes, version)?;
239        let u = U::from_bytes(bytes, version)?;
240        let v = V::from_bytes(bytes, version)?;
241        Ok((t, u, v))
242    }
243}
244
245impl<T> ChiaSerialize for Vec<T>
246where
247    T: ChiaSerialize,
248{
249    fn to_bytes(&self, version: ChiaProtocolVersion) -> Result<Vec<u8>, Error>
250    where
251        Self: Sized,
252    {
253        let mut bytes: Vec<u8> = Vec::new();
254        #[allow(clippy::cast_possible_truncation)]
255        bytes.extend((self.len() as u32).to_be_bytes());
256        for e in self {
257            bytes.extend(e.to_bytes(version)?);
258        }
259        Ok(bytes)
260    }
261    fn from_bytes<B: AsRef<[u8]>>(
262        bytes: &mut Cursor<B>,
263        version: ChiaProtocolVersion,
264    ) -> Result<Self, Error>
265    where
266        Self: Sized,
267    {
268        let mut u32_buf: [u8; 4] = [0; 4];
269        bytes.read_exact(&mut u32_buf)?;
270        let buf: Vec<T> = Vec::new();
271        let vec_len = u32::from_be_bytes(u32_buf);
272        if vec_len > 2048 {
273            warn!("Serializing Large Vec: {vec_len}");
274        }
275        (0..vec_len).try_fold(buf, |mut vec, _| {
276            vec.push(T::from_bytes(bytes, version)?);
277            Ok(vec)
278        })
279    }
280}
281
282macro_rules! impl_primitives {
283    ($($name: ident, $size:expr);*) => {
284        $(
285            impl ChiaSerialize for $name {
286                fn to_bytes(&self, _version: ChiaProtocolVersion) -> Result<Vec<u8>, Error> {
287                    Ok(self.to_be_bytes().to_vec())
288                }
289                fn from_bytes<T: AsRef<[u8]>>(bytes: &mut Cursor<T>, _version: ChiaProtocolVersion) -> Result<Self, std::io::Error> where Self: Sized,
290                {
291                    let remaining = bytes.get_ref().as_ref().len().saturating_sub(bytes.position() as usize);
292                    if remaining < $size {
293                        Err(Error::new(std::io::ErrorKind::InvalidInput, format!("Failed to Parse {}, expected length {}, found {}", stringify!($name), stringify!($size), remaining)))
294                    } else {
295                        let mut buffer: [u8; $size] = [0; $size];
296                        bytes.read_exact(&mut buffer)?;
297                        Ok($name::from_be_bytes(buffer))
298                    }
299                }
300            }
301        )*
302    };
303    ()=>{};
304}
305impl_primitives!(
306    i8, 1;
307    i16, 2;
308    i32, 4;
309    i64, 8;
310    i128, 16;
311    u8, 1;
312    u16, 2;
313    u32, 4;
314    u64, 8;
315    u128, 16;
316    f32, 4;
317    f64, 8
318);
319
320const MAX_DECODE_SIZE: u64 = 0x0004_0000_0000;
321
322#[allow(clippy::cast_possible_truncation)]
323pub fn encode_size(f: &mut dyn Write, size: u64) -> Result<(), Error> {
324    if size < 0x40 {
325        f.write_all(&[(0x80 | size) as u8])?;
326    } else if size < 0x2000 {
327        f.write_all(&[(0xc0 | (size >> 8)) as u8, ((size) & 0xff) as u8])?;
328    } else if size < 0x10_0000 {
329        f.write_all(&[
330            (0xe0 | (size >> 16)) as u8,
331            ((size >> 8) & 0xff) as u8,
332            ((size) & 0xff) as u8,
333        ])?;
334    } else if size < 0x800_0000 {
335        f.write_all(&[
336            (0xf0 | (size >> 24)) as u8,
337            ((size >> 16) & 0xff) as u8,
338            ((size >> 8) & 0xff) as u8,
339            ((size) & 0xff) as u8,
340        ])?;
341    } else if size < 0x4_0000_0000 {
342        f.write_all(&[
343            (0xf8 | (size >> 32)) as u8,
344            ((size >> 24) & 0xff) as u8,
345            ((size >> 16) & 0xff) as u8,
346            ((size >> 8) & 0xff) as u8,
347            ((size) & 0xff) as u8,
348        ])?;
349    } else {
350        return Err(Error::new(ErrorKind::InvalidData, "atom too big"));
351    }
352    Ok(())
353}
354
355pub fn decode_size(stream: &mut dyn Read, initial_b: u8) -> Result<u64, Error> {
356    if initial_b & 0x80 == 0 {
357        return Err(Error::new(ErrorKind::InvalidInput, "bad encoding"));
358    }
359    let mut bit_count = 0;
360    let mut bit_mask: u8 = 0x80;
361    let mut b = initial_b;
362    while b & bit_mask != 0 {
363        bit_count += 1;
364        b &= 0xff ^ bit_mask;
365        bit_mask >>= 1;
366    }
367    let mut size_blob: Vec<u8> = vec![0; bit_count];
368    size_blob[0] = b;
369    if bit_count > 1 {
370        stream.read_exact(&mut size_blob[1..])?;
371    }
372    let mut v = 0;
373    if size_blob.len() > 6 {
374        return Err(Error::new(ErrorKind::InvalidInput, "bad encoding"));
375    }
376    for b in &size_blob {
377        v <<= 8;
378        v += u64::from(*b);
379    }
380    if v >= MAX_DECODE_SIZE {
381        return Err(Error::new(ErrorKind::InvalidInput, "bad encoding"));
382    }
383    Ok(v)
384}
385
386impl<K: ChiaSerialize + Eq + Hash, V: ChiaSerialize> ChiaSerialize for HashMap<K, V> {
387    fn to_bytes(&self, version: ChiaProtocolVersion) -> Result<Vec<u8>, Error>
388    where
389        Self: Sized,
390    {
391        let mut bytes: Vec<u8> = Vec::new();
392        #[allow(clippy::cast_possible_truncation)]
393        bytes.extend((self.len() as u32).to_be_bytes());
394        for (k, v) in self {
395            bytes.extend(k.to_bytes(version)?);
396            bytes.extend(v.to_bytes(version)?);
397        }
398        Ok(bytes)
399    }
400
401    fn from_bytes<T: AsRef<[u8]>>(
402        bytes: &mut Cursor<T>,
403        version: ChiaProtocolVersion,
404    ) -> Result<Self, Error>
405    where
406        Self: Sized,
407    {
408        let mut u32_buf: [u8; 4] = [0; 4];
409        bytes.read_exact(&mut u32_buf)?;
410        let map_len = u32::from_be_bytes(u32_buf);
411        if map_len > 2048 {
412            warn!("Serializing Large Map: {map_len}");
413        }
414        let buf: HashMap<K, V> = HashMap::with_capacity(map_len as usize);
415        (0..map_len).try_fold(buf, |mut map, _| {
416            let key = K::from_bytes(bytes, version)?;
417            let value = V::from_bytes(bytes, version)?;
418            map.insert(key, value);
419            Ok(map)
420        })
421    }
422}