#![forbid(unsafe_code)]
use std::io::{Error, ErrorKind, Result};
use std::u64;
use byteorder::{BigEndian, ReadBytesExt};
use bytes::{Buf, BufMut, BytesMut, IntoBuf};
const HIGH_BIT: u8 = 0x80;
const LOWER_SEVEN_BITS: u8 = 0x7f;
pub fn predict_var_octet_string(length: usize) -> usize {
if length < 128 {
1 + length
} else {
let length_of_length = predict_var_uint_size(length as u64);
1 + length_of_length + length
}
}
fn predict_var_uint_size(value: u64) -> usize {
for i in 1..=8 {
let max = u64::MAX >> (64 - 8 * i);
if value <= max {
return i;
}
}
unreachable!()
}
pub fn extract_var_octet_string(mut buffer: BytesMut) -> Result<BytesMut> {
let buffer_length = buffer.len();
let mut reader = &buffer[..];
let content_length = reader.read_var_octet_string_length()?;
let content_offset = buffer_length - reader.len();
let mut remaining = buffer.split_off(content_offset);
if remaining.len() < content_length {
Err(Error::new(ErrorKind::UnexpectedEof, "buffer too small"))
} else {
Ok(remaining.split_to(content_length))
}
}
pub trait BufOerExt<'a> {
fn peek_var_octet_string(&self) -> Result<&'a [u8]>;
fn read_var_octet_string(&mut self) -> Result<&'a [u8]>;
fn skip(&mut self, discard_bytes: usize) -> Result<()>;
fn skip_var_octet_string(&mut self) -> Result<()>;
fn read_var_octet_string_length(&mut self) -> Result<usize>;
fn read_var_uint(&mut self) -> Result<u64>;
}
impl<'a> BufOerExt<'a> for &'a [u8] {
#[inline]
fn peek_var_octet_string(&self) -> Result<&'a [u8]> {
let mut peek = &self[..];
let actual_length = peek.read_var_octet_string_length()?;
let offset = self.len() - peek.len();
if peek.len() < actual_length {
Err(Error::new(ErrorKind::UnexpectedEof, "buffer too small"))
} else {
Ok(&self[offset..(offset + actual_length)])
}
}
#[inline]
fn read_var_octet_string(&mut self) -> Result<&'a [u8]> {
let actual_length = self.read_var_octet_string_length()?;
if self.len() < actual_length {
Err(Error::new(ErrorKind::UnexpectedEof, "buffer too small"))
} else {
let to_return = &self[..actual_length];
*self = &self[actual_length..];
Ok(to_return)
}
}
#[inline]
fn skip(&mut self, discard_bytes: usize) -> Result<()> {
if self.len() < discard_bytes {
Err(Error::new(ErrorKind::UnexpectedEof, "buffer too small"))
} else {
*self = &self[discard_bytes..];
Ok(())
}
}
#[inline]
fn skip_var_octet_string(&mut self) -> Result<()> {
let actual_length = self.read_var_octet_string_length()?;
self.skip(actual_length)
}
#[doc(hidden)]
#[inline]
fn read_var_octet_string_length(&mut self) -> Result<usize> {
let length = self.read_u8()?;
if length & HIGH_BIT != 0 {
let length_prefix_length = (length & LOWER_SEVEN_BITS) as usize;
if length_prefix_length > 8 {
Err(Error::new(
ErrorKind::InvalidData,
"length prefix too large",
))
} else {
Ok(self.read_uint::<BigEndian>(length_prefix_length)? as usize)
}
} else {
Ok(length as usize)
}
}
#[inline]
fn read_var_uint(&mut self) -> Result<u64> {
let size = self.read_var_octet_string_length()?;
if size == 0 {
Err(Error::new(ErrorKind::InvalidData, "zero-length VarUInt"))
} else if size > 8 {
Err(Error::new(ErrorKind::InvalidData, "VarUInt too large"))
} else {
Ok(self.read_uint::<BigEndian>(size)?)
}
}
}
pub trait MutBufOerExt: BufMut + Sized {
#[inline]
fn put_var_octet_string<B>(&mut self, buf: B)
where
B: IntoBuf,
{
let buf = buf.into_buf();
self.put_var_octet_string_length(buf.remaining());
self.put(buf);
}
#[doc(hidden)]
#[inline]
fn put_var_octet_string_length(&mut self, length: usize) {
if length < 128 {
self.put_u8(length as u8);
} else {
let length_of_length = predict_var_uint_size(length as u64);
self.put_u8(HIGH_BIT | length_of_length as u8);
self.put_uint_be(length as u64, length_of_length);
}
}
#[inline]
fn put_var_uint(&mut self, uint: u64) {
let size = predict_var_uint_size(uint);
self.put_var_octet_string_length(size);
self.put_uint_be(uint, size);
}
}
impl<B: BufMut + Sized> MutBufOerExt for B {}
#[cfg(test)]
mod test_functions {
use bytes::BytesMut;
use super::fixtures::*;
use super::*;
#[test]
fn test_predict_var_octet_string() {
let zeroes = &[0; 4096][..];
let mut buffer = BytesMut::with_capacity(5000);
for i in 0..4096 {
buffer.clear();
buffer.put_var_octet_string(&zeroes[..i]);
assert_eq!(
predict_var_octet_string(i),
buffer.len(),
"string_length={:?}",
i,
);
}
}
#[test]
fn test_predict_var_uint_size() {
assert_eq!(predict_var_uint_size(0), 1);
assert_eq!(predict_var_uint_size(1), 1);
assert_eq!(predict_var_uint_size(0xff), 1);
assert_eq!(predict_var_uint_size(0xff + 1), 2);
assert_eq!(predict_var_uint_size(u64::MAX - 1), 8);
assert_eq!(predict_var_uint_size(u64::MAX), 8);
}
#[test]
fn test_extract_var_octet_string() {
assert_eq!(
extract_var_octet_string(BytesMut::from(TWO_BYTE_VARSTR)).unwrap(),
BytesMut::from(&TWO_BYTE_VARSTR[1..3]),
);
assert_eq!(
extract_var_octet_string(BytesMut::new())
.unwrap_err()
.kind(),
ErrorKind::UnexpectedEof,
);
assert_eq!(
extract_var_octet_string(BytesMut::from(LENGTH_TOO_HIGH_VARSTR))
.unwrap_err()
.kind(),
ErrorKind::UnexpectedEof,
);
}
}
#[cfg(test)]
mod test_buf_oer_ext {
use lazy_static::lazy_static;
use super::fixtures::*;
use super::*;
lazy_static! {
static ref SIZE_128_VARSTR: Vec<u8> = {
let mut data = vec![0x81, 0x80];
data.extend(&[0x00; 128][..]);
data
};
static ref SIZE_5678_VARSTR: Vec<u8> = {
let mut data = vec![0x82, 0x16, 0x2e];
data.extend(&[0x00; 5678][..]);
data
};
}
#[test]
fn test_peek_var_octet_string() {
let tests: &[(Vec<u8>, &[u8])] = &[
(vec![0x00], &[]),
(ZERO_LENGTH_VARSTR.to_vec(), &[]),
(ONE_BYTE_VARSTR.to_vec(), &[0x01]),
(TWO_BYTE_VARSTR.to_vec(), &[0x01, 0x02]),
(SIZE_128_VARSTR.clone(), &[0; 128][..]),
(SIZE_5678_VARSTR.clone(), &[0; 5678][..]),
];
for (buffer, varstr) in tests {
assert_eq!((&buffer[..]).peek_var_octet_string().unwrap(), *varstr,);
}
assert_eq!(
(&[][..]).peek_var_octet_string().unwrap_err().kind(),
ErrorKind::UnexpectedEof,
);
assert_eq!(
LENGTH_TOO_HIGH_VARSTR
.peek_var_octet_string()
.unwrap_err()
.kind(),
ErrorKind::UnexpectedEof,
);
}
#[test]
fn test_skip() {
let mut empty = &[][..];
assert_eq!(empty.skip(1).unwrap_err().kind(), ErrorKind::UnexpectedEof,);
let mut reader = &[0x01, 0x02, 0x03][..];
assert!(reader.skip(2).is_ok());
assert_eq!(reader, &[0x03]);
}
#[test]
fn test_skip_var_octet_string() {
let tests: &[(Vec<u8>, usize)] = &[
(ZERO_LENGTH_VARSTR.to_vec(), 1),
(ONE_BYTE_VARSTR.to_vec(), 2),
(TWO_BYTE_VARSTR.to_vec(), 3),
(SIZE_128_VARSTR.clone(), SIZE_128_VARSTR.len()),
(SIZE_5678_VARSTR.clone(), SIZE_5678_VARSTR.len()),
];
for (buffer, offset) in tests {
let mut reader = &buffer[..];
assert!(reader.skip_var_octet_string().is_ok());
assert_eq!(reader.len(), buffer.len() - *offset);
}
assert_eq!(
(&[][..]).skip_var_octet_string().unwrap_err().kind(),
ErrorKind::UnexpectedEof,
);
assert_eq!(
(&LENGTH_TOO_HIGH_VARSTR[..])
.skip_var_octet_string()
.unwrap_err()
.kind(),
ErrorKind::UnexpectedEof,
);
}
#[test]
fn test_read_var_octet_string_length() {
let mut too_big: &[u8] = &[HIGH_BIT | 0x09];
assert_eq!(
too_big.read_var_octet_string_length().unwrap_err().kind(),
ErrorKind::InvalidData,
);
}
#[test]
fn test_read_var_uint() {
let tests: &[(Vec<u8>, u64, usize)] = &[
(vec![0x01, 0x00], 0, 2),
(vec![0x01, 0x09], 9, 2),
(vec![0x02, 0x01, 0x02], 0x0102, 3),
(vec![0x03, 0x01, 0x02, 0x03], 0x0001_0203, 4),
(vec![0x04, 0x01, 0x02, 0x03, 0x04], 0x0102_0304, 5),
(
vec![0x05, 0x01, 0x02, 0x03, 0x04, 0x05],
0x0001_0203_0405,
6,
),
(
vec![0x06, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06],
0x0102_0304_0506,
7,
),
(
vec![0x07, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07],
0x0001_0203_0405_0607,
8,
),
(
vec![0x08, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08],
0x0102_0304_0506_0708,
9,
),
];
for (buffer, value, offset) in tests {
let mut reader = &buffer[..];
assert_eq!(reader.read_var_uint().unwrap(), *value);
assert_eq!(reader.len(), buffer.len() - *offset);
}
let tests: &[(Vec<u8>, ErrorKind)] = &[
(vec![0x00], ErrorKind::InvalidData),
(vec![0x04], ErrorKind::UnexpectedEof),
(vec![0x04, 0x01, 0x02, 0x03], ErrorKind::UnexpectedEof),
(
vec![0x09, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09],
ErrorKind::InvalidData,
),
];
for (buffer, error_kind) in tests {
assert_eq!(
(&buffer[..]).read_var_uint().unwrap_err().kind(),
*error_kind,
);
}
}
}
#[cfg(test)]
mod buf_mut_oer_ext {
use super::*;
#[test]
fn test_put_var_octet_string() {
let long_varstr = &[0x00; 256][..];
let long_buffer = {
let mut buffer = vec![0x82, 0x01, 0x00];
buffer.extend_from_slice(&long_varstr);
buffer
};
let tests: &[(&[u8], &[u8])] = &[
(b"", b"\x00"),
(b"\xb0", b"\x01\xb0"),
(long_varstr, &long_buffer[..]),
];
for (varstr, buffer) in tests {
let mut writer = Vec::new();
writer.put_var_octet_string(varstr);
assert_eq!(&writer[..], *buffer);
}
}
#[test]
fn test_put_var_uint() {
let tests: &[(Vec<u8>, u64, usize)] = &[
(vec![0x01, 0x00], 0, 2),
(vec![0x01, 0x09], 9, 2),
(vec![0x02, 0x01, 0x02], 0x0102, 3),
(vec![0x03, 0x01, 0x02, 0x03], 0x0001_0203, 4),
(vec![0x04, 0x01, 0x02, 0x03, 0x04], 0x0102_0304, 5),
(
vec![0x05, 0x01, 0x02, 0x03, 0x04, 0x05],
0x0001_0203_0405,
6,
),
(
vec![0x06, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06],
0x0102_0304_0506,
7,
),
(
vec![0x07, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07],
0x0001_0203_0405_0607,
8,
),
(
vec![0x08, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08],
0x0102_0304_0506_0708,
9,
),
];
for (buffer, value, _offset) in tests {
let mut writer = Vec::new();
writer.put_var_uint(*value);
assert_eq!(writer, *buffer);
}
}
}
#[cfg(test)]
mod fixtures {
pub static ZERO_LENGTH_VARSTR: &[u8] = &[0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06];
pub static ONE_BYTE_VARSTR: &[u8] = &[0x01, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06];
pub static TWO_BYTE_VARSTR: &[u8] = &[0x02, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06];
pub static LENGTH_TOO_HIGH_VARSTR: &[u8] = &[0x07, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06];
}