Implement writers/readers

issue #62

Author: yuntai

src/util/ser.rs

@@ -0,0 +1,373 @@
+use std::result::Result;
+use std::io::{self, Read, Write};
+use std::fmt;
+use std::collections::HashMap;
+use std::hash::Hash;
+use std::mem;
+
+use secp256k1::Secp256k1;
+use secp256k1::key::PublicKey;
+use bitcoin::blockdata::script::Script;
+use bitcoin::util::hash::Sha256dHash;
+use std::marker::Sized;
+
+pub const MAX_VEC_SIZE: usize = 32 * 1024 * 1024;
+
+pub enum Error {
+	CorruptedData,
+	BadPublicKey,
+	IoError(io::Error),
+	InvalidLength,
+	InvalidValue,
+}
+
+impl fmt::Debug for Error {
+  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+		match *self {
+			Error::CorruptedData => f.write_str("corrupted data"),
+			Error::BadPublicKey => f.write_str("Bad public key"),
+			Error::IoError(ref e) => write!(f, "{}", e),
+			Error::InvalidLength => f.write_str("Invalid length"),
+			Error::InvalidValue => f.write_str("Invalid value"),
+		}
+  }
+}
+
+impl From<io::Error> for Error {
+  fn from(e: io::Error) -> Self {
+		Error::IoError(e)
+  }
+}
+
+pub struct Writer<W> { writer: W }
+pub struct Reader<R> { reader: R }
+
+pub trait Writeable<W: Write> {
+	fn write(&self, writer: &mut Writer<W>) -> Result<(), Error>;
+}
+
+pub trait Readable<R>
+	where Self: Sized,
+				R: Read
+{
+	fn read(reader: &mut Reader<R>) -> Result<Self, Error>;
+}
+
+macro_rules! writer_fn {
+	($name:ident, $val_type:ty, $convfn:ident) => {
+		#[inline]
+		fn $name(&mut self, v: $val_type) -> Result<(), Error> {
+			Ok(self.writer.write_all(&$convfn(v))?)
+		}
+	}
+}
+
+macro_rules! reader_fn {
+	($name:ident, $val_type:ty, $val_size: expr, $convfn:ident) => {
+		#[inline]
+		fn $name(&mut self) -> Result<$val_type, Error> {
+			let mut buf = [0; $val_size];
+			self.reader.read_exact(&mut buf)?;
+			Ok($convfn(&buf))
+		}
+	}
+}
+
+use util::byte_utils::{be64_to_array, be32_to_array, be16_to_array, slice_to_be16, slice_to_be32, slice_to_be64};
+
+impl<W: Write> Writer<W> {
+	pub fn new(writer: W) -> Writer<W> {
+		return Writer { writer }
+	}
+	pub fn into_inner(self) -> W { self.writer }
+	pub fn get_ref(&self) -> &W { &self.writer }
+	writer_fn!(write_u64, u64, be64_to_array);
+	writer_fn!(write_u32, u32, be32_to_array);
+	writer_fn!(write_u16, u16, be16_to_array);
+	fn write_u8(&mut self, v: u8) -> Result<(), Error> { 
+		Ok(self.writer.write_all(&[v])?)
+	}
+	fn write_bool(&mut self, v: bool) -> Result<(), Error> {
+		Ok(self.writer.write_all(&[if v {1} else {0}])?)
+	}
+}
+
+impl<R: Read> Reader<R> {
+	pub fn new(reader: R) -> Reader<R> {
+		return Reader { reader }
+	}
+  pub fn into_inner(self) -> R { self.reader }
+	pub fn get_ref(&self) -> &R { &self.reader }
+
+	reader_fn!(read_u16, u16, 2, slice_to_be16);
+	reader_fn!(read_u32, u32, 4, slice_to_be32);
+	reader_fn!(read_u64, u64, 8, slice_to_be64);
+
+	fn read_u8(&mut self) -> Result<u8, Error> { 
+		let mut buf = [0; 1];
+		self.reader.read_exact(&mut buf)?;
+		Ok(buf[0])
+	}
+	fn read_bool(&mut self) -> Result<bool, Error> {
+		let mut buf = [0; 1];
+		self.reader.read_exact(&mut buf)?;
+		if buf[0] != 0 && buf[0] != 1 {
+			return Err(Error::InvalidValue);
+		}
+		Ok(buf[0] == 1)
+	}
+	fn read_fixed_size(&mut self, sz: usize) -> Result<Vec<u8>, Error> {
+		let mut buf = vec![0; sz];
+		self.reader.read_exact(&mut buf)?;
+		Ok(buf)
+	}
+}
+
+macro_rules! impl_writeable_primitive {
+	($val_type:ty, $meth_write:ident, $meth_read:ident) => {
+		impl<W:Write> Writeable<W> for $val_type {
+			#[inline]
+			fn write(&self, writer: &mut Writer<W>) -> Result<(), Error> {
+				writer.$meth_write(*self)
+			}
+		}
+		impl<R:Read> Readable<R> for $val_type {
+			#[inline]
+			fn read(reader: &mut Reader<R>) -> Result<$val_type, Error> {
+				reader.$meth_read()
+			}
+		}
+	}
+}
+
+impl_writeable_primitive!(u64, write_u64, read_u64);
+impl_writeable_primitive!(u32, write_u32, read_u32);
+impl_writeable_primitive!(u16, write_u16, read_u16);
+impl_writeable_primitive!(u8, write_u8, read_u8);
+impl_writeable_primitive!(bool, write_bool, read_bool);
+
+// Arrays
+macro_rules! impl_array {
+	( $size:expr ) => (
+		impl<W, T> Writeable<W> for [T; $size]
+			where W: Write,
+						T: Writeable<W>,
+		{
+			#[inline]
+			fn write(&self, w: &mut Writer<W>) -> Result<(), Error> {
+					for i in self.iter() { i.write(w)?; }
+					Ok(())
+			}
+		}
+
+		impl<R, T> Readable<R> for [T; $size]
+			where R: Read,
+						T: Readable<R> + Copy,
+		{
+			#[inline]
+			fn read(r: &mut Reader<R>) -> Result<[T; $size], Error> {
+				let mut ret = [T::read(r)?; $size];
+				for item in ret.iter_mut().take($size).skip(1) { *item = T::read(r)?; }
+				Ok(ret)
+			}
+		}
+	);
+}
+
+impl_array!(32);
+impl_array!(33);
+
+// Tuples
+macro_rules! tuple_encode {
+	($($x:ident),*) => (
+		impl<W: Write, $($x: Writeable<W>),*> Writeable<W> for ($($x),*) {
+			#[inline]
+			#[allow(non_snake_case)]
+			fn write(&self, w: &mut Writer<W>) -> Result<(), Error> {
+					let &($(ref $x),*) = self;
+					$( $x.write(w)?; )*
+					Ok(())
+			}
+		}
+		impl<R: Read, $($x: Readable<R>),*> Readable<R> for ($($x),*) {
+			#[inline]
+			#[allow(non_snake_case)]
+			fn read(r: &mut Reader<R>) -> Result<($($x),*), Error> {
+					Ok(($({let $x = $x::read(r)?; $x }),*))
+			}
+		}
+	);
+}
+
+tuple_encode!(T0, T1);
+tuple_encode!(T0, T1, T2, T3);
+tuple_encode!(T0, T1, T2, T3, T4, T5);
+tuple_encode!(T0, T1, T2, T3, T4, T5, T6, T7);
+
+
+// References
+impl<W: Write, T: Writeable<W>> Writeable<W> for Box<T> {
+	#[inline]
+	fn write(&self, w: &mut Writer<W>) -> Result<(), Error> { (**self).write(w) }
+}
+
+impl<R: Read, T: Readable<R>> Readable<R> for Box<T> {
+	#[inline]
+	fn read(r: &mut Reader<R>) -> Result<Box<T>, Error> {
+		Ok(Box::new(<T as Readable<R>>::read(r)?))
+	}
+}
+
+// HashMap
+impl<W, K, V> Writeable<W> for HashMap<K, V>
+    where W: Write,
+          K: Writeable<W> + Eq + Hash,
+          V: Writeable<W>
+{
+    #[inline]
+    fn write(&self, w: &mut Writer<W>) -> Result<(), Error> {
+				(self.len() as u16).write(w)?;
+        for (key, value) in self.iter() {
+            key.write(w)?;
+            value.write(w)?;
+        }
+        Ok(())
+    }
+}
+
+impl<R, K, V> Readable<R> for HashMap<K, V>
+    where R: Read,
+          K: Readable<R> + Eq + Hash,
+          V: Readable<R>
+{
+	#[inline]
+	fn read(r: &mut Reader<R>) -> Result<HashMap<K, V>, Error> {
+		let len: u16 = Readable::read(r)?;
+		let mut ret = HashMap::with_capacity(len as usize);
+		for _ in 0..len {
+				ret.insert(K::read(r)?, V::read(r)?);
+		}
+		Ok(ret)
+	}
+}
+
+impl<W: Write, T: Writeable<W>> Writeable<W> for [T] {
+	#[inline]
+	fn write(&self, w: &mut Writer<W>) -> Result<(), Error> {
+		(self.len() as u16).write(w)?;
+		for c in self.iter() { c.write(w)?; }
+		Ok(())
+	}
+}
+
+// Vectors
+impl<W: Write, T: Writeable<W>> Writeable<W> for Vec<T> {
+	#[inline]
+  fn write(&self, w: &mut Writer<W>) -> Result<(), Error> { 
+		(&self[..]).write(w)
+	}
+}
+
+impl<R: Read, T: Readable<R>> Readable<R> for Vec<T> {
+	#[inline]
+	fn read(r: &mut Reader<R>) -> Result<Vec<T>, Error> {
+			let len: u16 = Readable::read(r)?;
+			let byte_size = (len as usize)
+												.checked_mul(mem::size_of::<T>())
+												.ok_or(Error::InvalidLength)?;
+			if byte_size > MAX_VEC_SIZE {
+					return Err(Error::InvalidLength);
+			}
+			let mut ret = Vec::with_capacity(len as usize);
+			for _ in 0..len { ret.push(T::read(r)?); }
+			Ok(ret)
+	}
+}
+
+impl<W: Write, T: Writeable<W>> Writeable<W> for Option<T> {
+    #[inline]
+    fn write(&self, w: &mut Writer<W>) -> Result<(), Error> {
+        match self {
+            Some(data) => {
+                true.write(w)?;
+                data.write(w)?;
+            }
+            None => { false.write(w)?; }
+        }
+        Ok(())
+    }
+}
+
+impl<R: Read, T:Readable<R>> Readable<R> for Option<T> {
+    #[inline]
+    fn read(r: &mut Reader<R>) -> Result<Option<T>, Error> {
+        let f: bool = Readable::read(r)?;
+        Ok(if f {
+            Some(Readable::read(r)?)
+        } else {
+            None
+        })
+    }
+}
+
+impl<W: Write> Writeable<W> for PublicKey {
+	fn write(&self, w: &mut Writer<W>) -> Result<(), Error> {
+		self.serialize().write(w)
+	}
+}
+
+impl<R: Read> Readable<R> for PublicKey {
+	fn read(r: &mut Reader<R>) -> Result<PublicKey, Error> {
+		let buf: [u8; 33] = Readable::read(r)?;
+    match PublicKey::from_slice(&Secp256k1::without_caps(), &buf) {
+      Ok(key) => Ok(key),
+      Err(_) => return Err(Error::BadPublicKey),
+    }
+	}
+}
+
+impl<W: Write> Writeable<W> for Sha256dHash {
+	fn write(&self, w: &mut Writer<W>) -> Result<(), Error> {
+		self.0.write(w);
+	}
+}
+
+impl<R: Read> Readable<R> for Sha256dHash {
+	fn read(r: &mut Reader<R>) -> Result<Sha256dHash, Error> {
+		let buf: [u8; 32] = Readable::read(r)?;
+		Ok(buf)
+	}
+}
+
+impl<W: Write> Writeable<W> for Script {
+	fn write(&self, w: &mut Writer<W>) -> Result<(), Error> {
+		self[..].write(w)
+	}
+}
+
+impl<R: Read> Readable<R> for Script {
+	fn read(r: &mut Reader<R>) -> Result<Script, Error> {
+		let sz: u16 = Readable::read(r)?;
+		let v = r.read_fixed_size(sz as usize)?;
+		Ok(Script::from(v))
+	}
+}
+
+macro_rules! impl_writeable {
+	($st:ident, {$($field:ident),*}) => {
+		impl<W: Write> Writeable<W> for $st {
+			fn write(&self, w: &mut Writer<W>) -> Result<(), Error> {
+				$( self.$field.write(w)?; )*
+				Ok(())
+			}
+		}
+
+		impl<R: Read> Readable<R> for $st {
+			fn read(r: &mut Reader<R>) -> Result<Self, Error> {
+				Ok(Self {
+					$($field: Readable::read(r)?),*
+				})
+			}
+		}
+	}
+}