Allow CompactSize deserialization to target any type that can be TryFrom<u64>

This commit is contained in:
Kris Nuttycombe 2021-08-30 13:58:23 -06:00
parent edcde252de
commit 427e6acbd4
2 changed files with 58 additions and 39 deletions

View File

@ -10,9 +10,10 @@
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use nonempty::NonEmpty;
use std::convert::TryFrom;
use std::io::{self, Read, Write};
const MAX_SIZE: usize = 0x02000000;
const MAX_SIZE: u64 = 0x02000000;
/// Namespace for functions for compact encoding of integers.
///
@ -22,17 +23,17 @@ pub struct CompactSize;
impl CompactSize {
/// Reads an integer encoded in compact form.
pub fn read<R: Read>(mut reader: R) -> io::Result<usize> {
pub fn read<R: Read>(mut reader: R) -> io::Result<u64> {
let flag = reader.read_u8()?;
match if flag < 253 {
Ok(flag as usize)
let result = if flag < 253 {
Ok(flag as u64)
} else if flag == 253 {
match reader.read_u16::<LittleEndian>()? {
n if n < 253 => Err(io::Error::new(
io::ErrorKind::InvalidInput,
"non-canonical CompactSize",
)),
n => Ok(n as usize),
n => Ok(n as u64),
}
} else if flag == 254 {
match reader.read_u32::<LittleEndian>()? {
@ -40,7 +41,7 @@ impl CompactSize {
io::ErrorKind::InvalidInput,
"non-canonical CompactSize",
)),
n => Ok(n as usize),
n => Ok(n as u64),
}
} else {
match reader.read_u64::<LittleEndian>()? {
@ -48,9 +49,11 @@ impl CompactSize {
io::ErrorKind::InvalidInput,
"non-canonical CompactSize",
)),
n => Ok(n as usize),
n => Ok(n),
}
}? {
}?;
match result {
s if s > MAX_SIZE => Err(io::Error::new(
io::ErrorKind::InvalidInput,
"CompactSize too large",
@ -59,6 +62,18 @@ impl CompactSize {
}
}
/// Reads an integer encoded in contact form and performs checked conversion
/// to the target type.
pub fn read_t<R: Read, T: TryFrom<u64>>(mut reader: R) -> io::Result<T> {
let n = Self::read(&mut reader)?;
<T>::try_from(n).map_err(|_| {
io::Error::new(
io::ErrorKind::InvalidInput,
"CompactSize value exceeds range of target type.",
)
})
}
/// Writes the provided `usize` value to the provided Writer in compact form.
pub fn write<W: Write>(mut writer: W, size: usize) -> io::Result<()> {
match size {
@ -92,7 +107,7 @@ impl Vector {
where
F: Fn(&mut R) -> io::Result<E>,
{
let count = CompactSize::read(&mut reader)?;
let count: usize = CompactSize::read_t(&mut reader)?;
Array::read(reader, count, func)
}
@ -193,34 +208,38 @@ impl Optional {
#[cfg(test)]
mod tests {
use super::*;
use std::convert::{TryFrom, TryInto};
use std::fmt::Debug;
#[test]
fn compact_size() {
macro_rules! eval {
($value:expr, $expected:expr) => {
let mut data = vec![];
CompactSize::write(&mut data, $value).unwrap();
assert_eq!(&data[..], &$expected[..]);
match CompactSize::read(&data[..]) {
Ok(n) => assert_eq!(n, $value),
Err(e) => panic!("Unexpected error: {:?}", e),
}
};
fn eval<T: TryFrom<u64> + TryInto<usize> + Eq + Debug + Copy>(value: T, expected: &[u8])
where
<T as TryInto<usize>>::Error: Debug,
{
let mut data = vec![];
CompactSize::write(&mut data, value.try_into().unwrap()).unwrap();
assert_eq!(&data[..], expected);
let result: io::Result<T> = CompactSize::read_t(&data[..]);
match result {
Ok(n) => assert_eq!(n, value),
Err(e) => panic!("Unexpected error: {:?}", e),
}
}
eval!(0, [0]);
eval!(1, [1]);
eval!(252, [252]);
eval!(253, [253, 253, 0]);
eval!(254, [253, 254, 0]);
eval!(255, [253, 255, 0]);
eval!(256, [253, 0, 1]);
eval!(256, [253, 0, 1]);
eval!(65535, [253, 255, 255]);
eval!(65536, [254, 0, 0, 1, 0]);
eval!(65537, [254, 1, 0, 1, 0]);
eval(0, &[0]);
eval(1, &[1]);
eval(252, &[252]);
eval(253, &[253, 253, 0]);
eval(254, &[253, 254, 0]);
eval(255, &[253, 255, 0]);
eval(256, &[253, 0, 1]);
eval(256, &[253, 0, 1]);
eval(65535, &[253, 255, 255]);
eval(65536, &[254, 0, 0, 1, 0]);
eval(65537, &[254, 1, 0, 1, 0]);
eval!(33554432, [254, 0, 0, 0, 2]);
eval(33554432, &[254, 0, 0, 0, 2]);
{
let value = 33554433;

View File

@ -117,15 +117,15 @@ impl TzeIn<<Authorized as Authorization>::Witness> {
pub fn read<R: Read>(mut reader: &mut R) -> io::Result<Self> {
let prevout = OutPoint::read(&mut reader)?;
let extension_id = CompactSize::read(&mut reader)?;
let mode = CompactSize::read(&mut reader)?;
let extension_id = CompactSize::read_t(&mut reader)?;
let mode = CompactSize::read_t(&mut reader)?;
let payload = Vector::read(&mut reader, |r| r.read_u8())?;
Ok(TzeIn {
prevout,
witness: tze::Witness {
extension_id: u32::try_from(extension_id).map_err(to_io_error)?,
mode: u32::try_from(mode).map_err(to_io_error)?,
extension_id,
mode,
payload: tze::AuthData(payload),
},
})
@ -159,15 +159,15 @@ impl TzeOut {
}
.map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "value out of range"))?;
let extension_id = CompactSize::read(&mut reader)?;
let mode = CompactSize::read(&mut reader)?;
let extension_id = CompactSize::read_t(&mut reader)?;
let mode = CompactSize::read_t(&mut reader)?;
let payload = Vector::read(&mut reader, |r| r.read_u8())?;
Ok(TzeOut {
value,
precondition: tze::Precondition {
extension_id: u32::try_from(extension_id).map_err(to_io_error)?,
mode: u32::try_from(mode).map_err(to_io_error)?,
extension_id,
mode,
payload,
},
})