use bytes::Bytes;
use futures::{Sink, SinkExt};
use httlib_huffman::DecoderSpeed;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite};
use crate::{
codec::{DecodeError, EncodeError, FixedLengthReader, StreamDecodeError, StreamEncodeError},
qpack::integer::{decode_integer, encode_integer},
};
pub async fn decode_string(
stream: impl AsyncRead,
prefix: u8,
n: u8,
) -> Result<(bool, Bytes), StreamDecodeError> {
tokio::pin!(stream);
let huffman = (prefix >> (n - 1)) & 1 == 1;
let length = decode_integer(stream.as_mut(), prefix, n - 1).await?;
let mut value = Vec::with_capacity((length as usize).min(8192));
FixedLengthReader::new(stream.as_mut(), length)
.read_to_end(&mut value)
.await?;
match huffman {
true => Ok((huffman, {
let mut decoded_value = vec![];
httlib_huffman::decode(&value, &mut decoded_value, DecoderSpeed::FourBits)
.map_err(DecodeError::from)?;
Bytes::from_owner(decoded_value)
})),
false => Ok((huffman, Bytes::from_owner(value))),
}
}
pub async fn encode_string<E>(
stream: impl AsyncWrite + Sink<Bytes, Error = E>,
mut prefix: u8,
n: u8,
huffman: bool,
data: Bytes,
) -> Result<(), StreamEncodeError>
where
StreamEncodeError: From<E>,
{
tokio::pin!(stream);
prefix |= (huffman as u8) << (n - 1);
match huffman {
true => {
let mut encoded_data = vec![];
httlib_huffman::encode(&data, &mut encoded_data).map_err(EncodeError::from)?;
encode_integer(stream.as_mut(), prefix, n - 1, encoded_data.len() as u64).await?;
stream.send(Bytes::from_owner(encoded_data)).await?;
Ok(())
}
false => {
encode_integer(stream.as_mut(), prefix, n - 1, data.len() as u64).await?;
stream.send(Bytes::from_owner(data)).await?;
Ok(())
}
}
}
#[cfg(test)]
mod tests {
use bytes::{Buf, Bytes};
use super::*;
use crate::buflist::BufList;
async fn round_trip_string(data: &[u8], huffman: bool) {
let mut writer = BufList::new();
let prefix: u8 = 0;
let n: u8 = 8;
encode_string(&mut writer, prefix, n, huffman, Bytes::from(data.to_vec()))
.await
.unwrap();
let total = writer.remaining();
let encoded = writer.copy_to_bytes(total);
let (decoded_huffman, decoded_data) =
decode_string(std::io::Cursor::new(&encoded[1..]), encoded[0], n)
.await
.unwrap();
assert_eq!(decoded_huffman, huffman);
assert_eq!(&decoded_data[..], data);
}
#[tokio::test]
async fn plain_empty_string() {
round_trip_string(b"", false).await;
}
#[tokio::test]
async fn plain_ascii_string() {
round_trip_string(b"hello", false).await;
}
#[tokio::test]
async fn plain_longer_string() {
round_trip_string(b"www.example.com", false).await;
}
#[tokio::test]
async fn huffman_ascii_string() {
round_trip_string(b"hello", true).await;
}
#[tokio::test]
async fn huffman_longer_string() {
round_trip_string(b"www.example.com", true).await;
}
#[tokio::test]
async fn huffman_empty_string() {
round_trip_string(b"", true).await;
}
#[tokio::test]
async fn huffman_flag_bit_set_correctly() {
let mut writer = BufList::new();
let n: u8 = 8;
encode_string(&mut writer, 0, n, false, Bytes::from_static(b"a"))
.await
.unwrap();
let plain_byte = Buf::chunk(&writer)[0];
assert_eq!(plain_byte & 0x80, 0, "H bit should be 0 for plain");
let mut writer2 = BufList::new();
encode_string(&mut writer2, 0, n, true, Bytes::from_static(b"a"))
.await
.unwrap();
let huff_byte = Buf::chunk(&writer2)[0];
assert_eq!(huff_byte & 0x80, 0x80, "H bit should be 1 for huffman");
}
mod proptest_roundtrip {
use bytes::{Buf, Bytes};
use proptest::prelude::*;
use super::*;
use crate::buflist::BufList;
proptest! {
#[test]
fn qpack_string_plain_roundtrip(data in proptest::collection::vec(any::<u8>(), 0..256)) {
let rt = tokio::runtime::Builder::new_current_thread().build().unwrap();
rt.block_on(async {
let mut writer = BufList::new();
encode_string(&mut writer, 0, 8, false, Bytes::from(data.clone()))
.await
.unwrap();
let total = writer.remaining();
let encoded = writer.copy_to_bytes(total);
let (huffman, decoded) =
decode_string(std::io::Cursor::new(&encoded[1..]), encoded[0], 8)
.await
.unwrap();
prop_assert!(!huffman);
prop_assert_eq!(&decoded[..], &data[..]);
Ok(())
})?;
}
#[test]
fn qpack_string_huffman_roundtrip(data in proptest::collection::vec(0u8..128, 0..256)) {
let rt = tokio::runtime::Builder::new_current_thread().build().unwrap();
rt.block_on(async {
let mut writer = BufList::new();
encode_string(&mut writer, 0, 8, true, Bytes::from(data.clone()))
.await
.unwrap();
let total = writer.remaining();
let encoded = writer.copy_to_bytes(total);
let (huffman, decoded) =
decode_string(std::io::Cursor::new(&encoded[1..]), encoded[0], 8)
.await
.unwrap();
prop_assert!(huffman);
prop_assert_eq!(&decoded[..], &data[..]);
Ok(())
})?;
}
}
}
}