112 lines
3.4 KiB
Rust
112 lines
3.4 KiB
Rust
use std::{
|
|
convert::TryInto,
|
|
io,
|
|
net::{IpAddr, SocketAddr},
|
|
};
|
|
|
|
use byteorder::{BigEndian, LittleEndian, WriteBytesExt};
|
|
|
|
use super::MAX_PROTOCOL_MESSAGE_LEN;
|
|
|
|
/// Extends [`Write`] with methods for writing Zcash/Bitcoin types.
|
|
///
|
|
/// [`Write`]: https://doc.rust-lang.org/std/io/trait.Write.html
|
|
pub trait WriteZcashExt: io::Write {
|
|
/// Writes a `u64` using the Bitcoin `CompactSize` encoding.
|
|
///
|
|
/// # Panics
|
|
///
|
|
/// Zebra panics on sizes greater than the protocol message length limit.
|
|
/// This is a defence-in-depth for memory preallocation attacks.
|
|
/// (These sizes should be impossible, because each array item takes at
|
|
/// least one byte.)
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use zebra_chain::serialization::WriteZcashExt;
|
|
///
|
|
/// let mut buf = Vec::new();
|
|
/// buf.write_compactsize(0x12).unwrap();
|
|
/// assert_eq!(buf, b"\x12");
|
|
///
|
|
/// let mut buf = Vec::new();
|
|
/// buf.write_compactsize(0xfd).unwrap();
|
|
/// assert_eq!(buf, b"\xfd\xfd\x00");
|
|
///
|
|
/// let mut buf = Vec::new();
|
|
/// buf.write_compactsize(0xaafd).unwrap();
|
|
/// assert_eq!(buf, b"\xfd\xfd\xaa");
|
|
/// ```
|
|
///
|
|
/// Sizes greater than the maximum network message length are invalid
|
|
/// and cause a panic:
|
|
/// ```should_panic
|
|
/// # use zebra_chain::serialization::WriteZcashExt;
|
|
/// let mut buf = Vec::new();
|
|
/// buf.write_compactsize(0xbbaafd).unwrap_err();
|
|
/// buf.write_compactsize(0x22ccbbaafd).unwrap_err();
|
|
/// ```
|
|
#[inline]
|
|
fn write_compactsize(&mut self, n: u64) -> io::Result<()> {
|
|
// # Security
|
|
// Defence-in-depth for memory DoS via preallocation.
|
|
//
|
|
assert!(
|
|
n <= MAX_PROTOCOL_MESSAGE_LEN
|
|
.try_into()
|
|
.expect("usize fits in u64"),
|
|
"CompactSize larger than protocol message limit"
|
|
);
|
|
|
|
match n {
|
|
0x0000_0000..=0x0000_00fc => self.write_u8(n as u8),
|
|
0x0000_00fd..=0x0000_ffff => {
|
|
self.write_u8(0xfd)?;
|
|
self.write_u16::<LittleEndian>(n as u16)
|
|
}
|
|
0x0001_0000..=0xffff_ffff => {
|
|
self.write_u8(0xfe)?;
|
|
self.write_u32::<LittleEndian>(n as u32)
|
|
}
|
|
_ => {
|
|
self.write_u8(0xff)?;
|
|
self.write_u64::<LittleEndian>(n)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Write an `IpAddr` in Bitcoin format.
|
|
#[inline]
|
|
fn write_ip_addr(&mut self, addr: IpAddr) -> io::Result<()> {
|
|
use std::net::IpAddr::*;
|
|
let v6_addr = match addr {
|
|
V4(ref v4) => v4.to_ipv6_mapped(),
|
|
V6(v6) => v6,
|
|
};
|
|
self.write_all(&v6_addr.octets())
|
|
}
|
|
|
|
/// Write a `SocketAddr` in Bitcoin format.
|
|
#[inline]
|
|
fn write_socket_addr(&mut self, addr: SocketAddr) -> io::Result<()> {
|
|
self.write_ip_addr(addr.ip())?;
|
|
self.write_u16::<BigEndian>(addr.port())
|
|
}
|
|
|
|
/// Convenience method to write exactly 32 u8's.
|
|
#[inline]
|
|
fn write_32_bytes(&mut self, bytes: &[u8; 32]) -> io::Result<()> {
|
|
self.write_all(bytes)
|
|
}
|
|
|
|
/// Convenience method to write exactly 64 u8's.
|
|
#[inline]
|
|
fn write_64_bytes(&mut self, bytes: &[u8; 64]) -> io::Result<()> {
|
|
self.write_all(bytes)
|
|
}
|
|
}
|
|
|
|
/// Mark all types implementing `Write` as implementing the extension.
|
|
impl<W: io::Write + ?Sized> WriteZcashExt for W {}
|