Revision control
Copy as Markdown
Other Tools
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// option. This file may not be copied, modified, or distributed
// except according to those terms.
use std::{mem, str};
use neqo_common::{qdebug, qerror};
use neqo_transport::{Connection, StreamId};
use crate::{huffman::decode_huffman, prefix::Prefix, Error, Res};
pub trait ReadByte {
/// # Errors
///
/// Return error occurred while reading a byte.
/// The exact error depends on trait implementation.
fn read_byte(&mut self) -> Res<u8>;
}
pub trait Reader {
/// # Errors
///
/// Return error occurred while reading date into a buffer.
/// The exact error depends on trait implementation.
fn read(&mut self, buf: &mut [u8]) -> Res<usize>;
}
pub(crate) struct ReceiverConnWrapper<'a> {
conn: &'a mut Connection,
stream_id: StreamId,
}
impl ReadByte for ReceiverConnWrapper<'_> {
fn read_byte(&mut self) -> Res<u8> {
let mut b = [0];
match self.conn.stream_recv(self.stream_id, &mut b)? {
(_, true) => Err(Error::ClosedCriticalStream),
(0, false) => Err(Error::NeedMoreData),
_ => Ok(b[0]),
}
}
}
impl Reader for ReceiverConnWrapper<'_> {
fn read(&mut self, buf: &mut [u8]) -> Res<usize> {
match self.conn.stream_recv(self.stream_id, buf)? {
(_, true) => Err(Error::ClosedCriticalStream),
(amount, false) => Ok(amount),
}
}
}
impl<'a> ReceiverConnWrapper<'a> {
pub fn new(conn: &'a mut Connection, stream_id: StreamId) -> Self {
Self { conn, stream_id }
}
}
/// This is only used by header decoder therefore all errors are `DecompressionFailed`.
/// A header block is read entirely before decoding it, therefore if there is not enough
/// data in the buffer an error `DecompressionFailed` will be return.
pub(crate) struct ReceiverBufferWrapper<'a> {
buf: &'a [u8],
offset: usize,
}
impl ReadByte for ReceiverBufferWrapper<'_> {
fn read_byte(&mut self) -> Res<u8> {
if self.offset == self.buf.len() {
Err(Error::DecompressionFailed)
} else {
let b = self.buf[self.offset];
self.offset += 1;
Ok(b)
}
}
}
impl<'a> ReceiverBufferWrapper<'a> {
pub const fn new(buf: &'a [u8]) -> Self {
Self { buf, offset: 0 }
}
pub const fn peek(&self) -> Res<u8> {
if self.offset == self.buf.len() {
Err(Error::DecompressionFailed)
} else {
Ok(self.buf[self.offset])
}
}
pub const fn done(&self) -> bool {
self.offset == self.buf.len()
}
/// The function decodes varint with a prefixed, i.e. ignores `prefix_len` bits of the first
/// byte.
/// `ReceiverBufferWrapper` is only used for decoding header blocks. The header blocks are read
/// entirely before a decoding starts, therefore any incomplete varint because of reaching the
/// end of a buffer will be treated as the `DecompressionFailed` error.
pub fn read_prefixed_int(&mut self, prefix_len: u8) -> Res<u64> {
debug_assert!(prefix_len < 8);
let first_byte = self.read_byte()?;
let mut reader = IntReader::new(first_byte, prefix_len);
reader.read(self)
}
/// Do not use `LiteralReader` here to avoid copying data.
/// The function decoded a literal with a prefix:
/// 1) ignores `prefix_len` bits of the first byte,
/// 2) reads "huffman bit"
/// 3) decode varint that is the length of a literal
/// 4) reads the literal
/// 5) performs huffman decoding if needed.
///
/// `ReceiverBufferWrapper` is only used for decoding header blocks. The header blocks are read
/// entirely before a decoding starts, therefore any incomplete varint or literal because of
/// reaching the end of a buffer will be treated as the `DecompressionFailed` error.
pub fn read_literal_from_buffer(&mut self, prefix_len: u8) -> Res<String> {
debug_assert!(prefix_len < 7);
let first_byte = self.read_byte()?;
let use_huffman = (first_byte & (0x80 >> prefix_len)) != 0;
let mut int_reader = IntReader::new(first_byte, prefix_len + 1);
let length: usize = int_reader
.read(self)?
.try_into()
.or(Err(Error::DecompressionFailed))?;
if use_huffman {
Ok(parse_utf8(&decode_huffman(self.slice(length)?)?)?.to_string())
} else {
Ok(parse_utf8(self.slice(length)?)?.to_string())
}
}
fn slice(&mut self, len: usize) -> Res<&[u8]> {
if self.offset + len > self.buf.len() {
Err(Error::DecompressionFailed)
} else {
let start = self.offset;
self.offset += len;
Ok(&self.buf[start..self.offset])
}
}
}
/// This is varint reader that can take into account a prefix.
#[derive(Debug)]
pub struct IntReader {
value: u64,
cnt: u8,
done: bool,
}
impl IntReader {
/// `IntReader` is created by suppling the first byte anf prefix length.
/// A varint may take only one byte, In that case already the first by has set state to done.
///
/// # Panics
///
/// When `prefix_len` is 8 or larger.
#[must_use]
pub fn new(first_byte: u8, prefix_len: u8) -> Self {
debug_assert!(prefix_len < 8, "prefix cannot larger than 7.");
let mask = if prefix_len == 0 {
0xff
} else {
(1 << (8 - prefix_len)) - 1
};
let value = u64::from(first_byte & mask);
Self {
value,
cnt: 0,
done: value < u64::from(mask),
}
}
/// # Panics
///
/// Never, but rust doesn't know that.
#[must_use]
pub fn make(first_byte: u8, prefixes: &[Prefix]) -> Self {
for prefix in prefixes {
if prefix.cmp_prefix(first_byte) {
return Self::new(first_byte, prefix.len());
}
}
unreachable!();
}
/// This function reads bytes until the varint is decoded or until stream/buffer does not
/// have any more date.
///
/// # Errors
///
/// Possible errors are:
/// 1) `NeedMoreData` if the reader needs more data,
/// 2) `IntegerOverflow`,
/// 3) Any `ReadByte`'s error
pub fn read<R: ReadByte>(&mut self, s: &mut R) -> Res<u64> {
let mut b: u8;
while !self.done {
b = s.read_byte()?;
if (self.cnt == 63) && (b > 1 || (b == 1 && ((self.value >> 63) == 1))) {
qerror!("Error decoding prefixed encoded int - IntegerOverflow");
return Err(Error::IntegerOverflow);
}
self.value += u64::from(b & 0x7f) << self.cnt;
if (b & 0x80) == 0 {
self.done = true;
}
self.cnt += 7;
if self.cnt >= 64 {
self.done = true;
}
}
Ok(self.value)
}
}
#[derive(Debug, Default)]
enum LiteralReaderState {
#[default]
ReadHuffman,
ReadLength {
reader: IntReader,
},
ReadLiteral {
offset: usize,
},
Done,
}
/// This is decoder of a literal with a prefix:
/// 1) ignores `prefix_len` bits of the first byte,
/// 2) reads "huffman bit"
/// 3) decode varint that is the length of a literal
/// 4) reads the literal
/// 5) performs huffman decoding if needed.
#[derive(Debug, Default)]
pub struct LiteralReader {
state: LiteralReaderState,
literal: Vec<u8>,
use_huffman: bool,
}
impl LiteralReader {
/// Creates `LiteralReader` with the first byte. This constructor is always used
/// when a litreral has a prefix.
/// For literals without a prefix please use the default constructor.
///
/// # Panics
///
/// If `prefix_len` is 8 or more.
#[must_use]
pub fn new_with_first_byte(first_byte: u8, prefix_len: u8) -> Self {
assert!(prefix_len < 8);
Self {
state: LiteralReaderState::ReadLength {
reader: IntReader::new(first_byte, prefix_len + 1),
},
literal: Vec::new(),
use_huffman: (first_byte & (0x80 >> prefix_len)) != 0,
}
}
/// This function reads bytes until the literal is decoded or until stream/buffer does not
/// have any more date ready.
///
/// # Errors
///
/// Possible errors are:
/// 1) `NeedMoreData` if the reader needs more data,
/// 2) `IntegerOverflow`
/// 3) Any `ReadByte`'s error
///
/// It returns value if reading the literal is done or None if it needs more data.
///
/// # Panics
///
/// When this object is complete.
pub fn read<T: ReadByte + Reader>(&mut self, s: &mut T) -> Res<Vec<u8>> {
loop {
qdebug!("state = {:?}", self.state);
match &mut self.state {
LiteralReaderState::ReadHuffman => {
let b = s.read_byte()?;
self.use_huffman = (b & 0x80) != 0;
self.state = LiteralReaderState::ReadLength {
reader: IntReader::new(b, 1),
};
}
LiteralReaderState::ReadLength { reader } => {
let v = reader.read(s)?;
self.literal
.resize(v.try_into().or(Err(Error::Decoding))?, 0x0);
self.state = LiteralReaderState::ReadLiteral { offset: 0 };
}
LiteralReaderState::ReadLiteral { offset } => {
let amount = s.read(&mut self.literal[*offset..])?;
*offset += amount;
if *offset == self.literal.len() {
self.state = LiteralReaderState::Done;
if self.use_huffman {
break Ok(decode_huffman(&self.literal)?);
}
break Ok(mem::take(&mut self.literal));
}
break Err(Error::NeedMoreData);
}
LiteralReaderState::Done => {
panic!("Should not call read() in this state.");
}
}
}
}
}
/// This is a helper function used only by `ReceiverBufferWrapper`, therefore it returns
/// `DecompressionFailed` if any error happens.
///
/// # Errors
///
/// If an parsing error occurred, the function returns `BadUtf8`.
pub fn parse_utf8(v: &[u8]) -> Res<&str> {
str::from_utf8(v).map_err(|_| Error::BadUtf8)
}
#[cfg(test)]
pub(crate) mod test_receiver {
use std::collections::VecDeque;
use super::{Error, ReadByte, Reader, Res};
#[derive(Default)]
pub struct TestReceiver {
buf: VecDeque<u8>,
}
impl ReadByte for TestReceiver {
fn read_byte(&mut self) -> Res<u8> {
self.buf.pop_back().ok_or(Error::NeedMoreData)
}
}
impl Reader for TestReceiver {
fn read(&mut self, buf: &mut [u8]) -> Res<usize> {
let len = if buf.len() > self.buf.len() {
self.buf.len()
} else {
buf.len()
};
for item in buf.iter_mut().take(len) {
*item = self.buf.pop_back().ok_or(Error::NeedMoreData)?;
}
Ok(len)
}
}
impl TestReceiver {
pub fn write(&mut self, buf: &[u8]) {
for b in buf {
self.buf.push_front(*b);
}
}
}
}
#[cfg(test)]
mod tests {
use test_receiver::TestReceiver;
use super::{
parse_utf8, str, test_receiver, Error, IntReader, LiteralReader, ReadByte,
ReceiverBufferWrapper, Res,
};
const TEST_CASES_NUMBERS: [(&[u8], u8, u64); 7] = [
(&[0xEA], 3, 10),
(&[0x0A], 3, 10),
(&[0x8A], 3, 10),
(&[0xFF, 0x9A, 0x0A], 3, 1337),
(&[0x1F, 0x9A, 0x0A], 3, 1337),
(&[0x9F, 0x9A, 0x0A], 3, 1337),
(&[0x2A], 0, 42),
];
#[test]
fn read_prefixed_int() {
for (buf, prefix_len, value) in &TEST_CASES_NUMBERS {
let mut reader = IntReader::new(buf[0], *prefix_len);
let mut test_receiver: TestReceiver = TestReceiver::default();
test_receiver.write(&buf[1..]);
assert_eq!(reader.read(&mut test_receiver), Ok(*value));
}
}
#[test]
fn read_prefixed_int_with_more_data_in_buffer() {
for (buf, prefix_len, value) in &TEST_CASES_NUMBERS {
let mut reader = IntReader::new(buf[0], *prefix_len);
let mut test_receiver: TestReceiver = TestReceiver::default();
test_receiver.write(&buf[1..]);
// add some more data
test_receiver.write(&[0x0, 0x0, 0x0]);
assert_eq!(reader.read(&mut test_receiver), Ok(*value));
}
}
#[test]
fn read_prefixed_int_slow_writer() {
let (buf, prefix_len, value) = &TEST_CASES_NUMBERS[4];
let mut reader = IntReader::new(buf[0], *prefix_len);
let mut test_receiver: TestReceiver = TestReceiver::default();
// data has not been received yet, reading IntReader will return Err(Error::NeedMoreData).
assert_eq!(reader.read(&mut test_receiver), Err(Error::NeedMoreData));
// Write one byte.
test_receiver.write(&buf[1..2]);
// data has not been received yet, reading IntReader will return Err(Error::NeedMoreData).
assert_eq!(reader.read(&mut test_receiver), Err(Error::NeedMoreData));
// Write one byte.
test_receiver.write(&buf[2..]);
// Now prefixed int is complete.
assert_eq!(reader.read(&mut test_receiver), Ok(*value));
}
type TestSetup = (&'static [u8], u8, Res<u64>);
const TEST_CASES_BIG_NUMBERS: [TestSetup; 3] = [
(
&[
0xFF, 0x80, 0xFE, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x01,
],
0,
Ok(0xFFFF_FFFF_FFFF_FFFF),
),
(
&[
0xFF, 0x81, 0xFE, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x01,
],
0,
Err(Error::IntegerOverflow),
),
(
&[
0xFF, 0x80, 0xFE, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x02,
],
0,
Err(Error::IntegerOverflow),
),
];
#[test]
fn read_prefixed_int_big_number() {
for (buf, prefix_len, value) in &TEST_CASES_BIG_NUMBERS {
let mut reader = IntReader::new(buf[0], *prefix_len);
let mut test_receiver: TestReceiver = TestReceiver::default();
test_receiver.write(&buf[1..]);
assert_eq!(reader.read(&mut test_receiver), *value);
}
}
const TEST_CASES_LITERAL: [(&[u8], u8, &str); 9] = [
// No Huffman
(
&[
0x0a, 0x63, 0x75, 0x73, 0x74, 0x6f, 0x6d, 0x2d, 0x6b, 0x65, 0x79,
],
1,
"custom-key",
),
(
&[
0x0a, 0x63, 0x75, 0x73, 0x74, 0x6f, 0x6d, 0x2d, 0x6b, 0x65, 0x79,
],
3,
"custom-key",
),
(
&[
0xea, 0x63, 0x75, 0x73, 0x74, 0x6f, 0x6d, 0x2d, 0x6b, 0x65, 0x79,
],
3,
"custom-key",
),
(
&[
0x0d, 0x63, 0x75, 0x73, 0x74, 0x6f, 0x6d, 0x2d, 0x68, 0x65, 0x61, 0x64, 0x65, 0x72,
],
1,
"custom-header",
),
// With Huffman
(&[0x15, 0xae, 0xc3, 0x77, 0x1a, 0x4b], 3, "private"),
(
&[
0x56, 0xd0, 0x7a, 0xbe, 0x94, 0x10, 0x54, 0xd4, 0x44, 0xa8, 0x20, 0x05, 0x95, 0x04,
0x0b, 0x81, 0x66, 0xe0, 0x82, 0xa6, 0x2d, 0x1b, 0xff,
],
1,
"Mon, 21 Oct 2013 20:13:21 GMT",
),
(
&[
0xff, 0x0f, 0xd0, 0x7a, 0xbe, 0x94, 0x10, 0x54, 0xd4, 0x44, 0xa8, 0x20, 0x05, 0x95,
0x04, 0x0b, 0x81, 0x66, 0xe0, 0x82, 0xa6, 0x2d, 0x1b, 0xff,
],
4,
"Mon, 21 Oct 2013 20:13:21 GMT",
),
(
&[
0x51, 0x9d, 0x29, 0xad, 0x17, 0x18, 0x63, 0xc7, 0x8f, 0x0b, 0x97, 0xc8, 0xe9, 0xae,
0x82, 0xae, 0x43, 0xd3,
],
1,
),
(
&[
0x91, 0x9d, 0x29, 0xad, 0x17, 0x18, 0x63, 0xc7, 0x8f, 0x0b, 0x97, 0xc8, 0xe9, 0xae,
0x82, 0xae, 0x43, 0xd3,
],
0,
),
];
#[test]
fn read_literal() {
for (buf, prefix_len, value) in &TEST_CASES_LITERAL {
let mut reader = LiteralReader::new_with_first_byte(buf[0], *prefix_len);
let mut test_receiver: TestReceiver = TestReceiver::default();
test_receiver.write(&buf[1..]);
assert_eq!(
parse_utf8(&reader.read(&mut test_receiver).unwrap()).unwrap(),
*value
);
}
}
#[test]
fn read_prefixed_int_receiver_buffer_wrapper() {
for (buf, prefix_len, value) in &TEST_CASES_NUMBERS {
let mut buffer = ReceiverBufferWrapper::new(buf);
let mut reader = IntReader::new(buffer.read_byte().unwrap(), *prefix_len);
assert_eq!(reader.read(&mut buffer), Ok(*value));
}
}
#[test]
fn read_prefixed_int_big_receiver_buffer_wrapper() {
for (buf, prefix_len, value) in &TEST_CASES_BIG_NUMBERS {
let mut buffer = ReceiverBufferWrapper::new(buf);
let mut reader = IntReader::new(buffer.read_byte().unwrap(), *prefix_len);
assert_eq!(reader.read(&mut buffer), *value);
}
}
#[test]
fn read_literal_receiver_buffer_wrapper() {
for (buf, prefix_len, value) in &TEST_CASES_LITERAL {
let mut buffer = ReceiverBufferWrapper::new(buf);
assert_eq!(
buffer.read_literal_from_buffer(*prefix_len).unwrap(),
*value
);
}
}
#[test]
fn read_failure_receiver_buffer_wrapper_number() {
let (buf, prefix_len, _) = &TEST_CASES_NUMBERS[4];
let mut buffer = ReceiverBufferWrapper::new(&buf[..1]);
let mut reader = IntReader::new(buffer.read_byte().unwrap(), *prefix_len);
assert_eq!(reader.read(&mut buffer), Err(Error::DecompressionFailed));
}
#[test]
fn read_failure_receiver_buffer_wrapper_literal() {
let (buf, prefix_len, _) = &TEST_CASES_LITERAL[0];
let mut buffer = ReceiverBufferWrapper::new(&buf[..6]);
assert_eq!(
buffer.read_literal_from_buffer(*prefix_len),
Err(Error::DecompressionFailed)
);
}
}