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, Chia0_0_35 = 35, Chia0_0_36 = 36, #[default]
27 Chia0_0_37 = 37, }
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}