use std::future::poll_fn;
use std::io::{self, Cursor};
use std::ops::RangeInclusive;
use std::pin::Pin;
use std::task::{ready, Context, Poll};
use bytes::{Buf, BufMut, Bytes, BytesMut};
use pin_project_lite::pin_project;
use tokio::io::{AsyncBufRead, AsyncRead, AsyncReadExt, ReadBuf};
use crate::wire::{ProtocolVersion, EMPTY_BYTES};
use super::{Error, NixRead};
pub struct NixReaderBuilder {
buf: Option<BytesMut>,
reserved_buf_size: usize,
max_buf_size: usize,
version: ProtocolVersion,
}
impl Default for NixReaderBuilder {
fn default() -> Self {
Self {
buf: Default::default(),
reserved_buf_size: 8192,
max_buf_size: 8192,
version: Default::default(),
}
}
}
impl NixReaderBuilder {
pub fn set_buffer(mut self, buf: BytesMut) -> Self {
self.buf = Some(buf);
self
}
pub fn set_reserved_buf_size(mut self, size: usize) -> Self {
self.reserved_buf_size = size;
self
}
pub fn set_max_buf_size(mut self, size: usize) -> Self {
self.max_buf_size = size;
self
}
pub fn set_version(mut self, version: ProtocolVersion) -> Self {
self.version = version;
self
}
pub fn build<R>(self, reader: R) -> NixReader<R> {
let buf = self.buf.unwrap_or_else(|| BytesMut::with_capacity(0));
NixReader {
buf,
inner: reader,
reserved_buf_size: self.reserved_buf_size,
max_buf_size: self.max_buf_size,
version: self.version,
}
}
}
pin_project! {
pub struct NixReader<R> {
#[pin]
inner: R,
buf: BytesMut,
reserved_buf_size: usize,
max_buf_size: usize,
version: ProtocolVersion,
}
}
impl NixReader<Cursor<Vec<u8>>> {
pub fn builder() -> NixReaderBuilder {
NixReaderBuilder::default()
}
}
impl<R> NixReader<R>
where
R: AsyncReadExt,
{
pub fn new(reader: R) -> NixReader<R> {
NixReader::builder().build(reader)
}
pub fn buffer(&self) -> &[u8] {
&self.buf[..]
}
#[cfg(test)]
pub(crate) fn buffer_mut(&mut self) -> &mut BytesMut {
&mut self.buf
}
/// Remaining capacity in internal buffer
pub fn remaining_mut(&self) -> usize {
self.buf.capacity() - self.buf.len()
}
fn poll_force_fill_buf(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<io::Result<usize>> {
// Ensure that buffer has space for at least reserved_buf_size bytes
if self.remaining_mut() < self.reserved_buf_size {
let me = self.as_mut().project();
me.buf.reserve(*me.reserved_buf_size);
}
let me = self.project();
let n = {
let dst = me.buf.spare_capacity_mut();
let mut buf = ReadBuf::uninit(dst);
let ptr = buf.filled().as_ptr();
ready!(me.inner.poll_read(cx, &mut buf)?);
// Ensure the pointer does not change from under us
assert_eq!(ptr, buf.filled().as_ptr());
buf.filled().len()
};
// SAFETY: This is guaranteed to be the number of initialized (and read)
// bytes due to the invariants provided by `ReadBuf::filled`.
unsafe {
me.buf.advance_mut(n);
}
Poll::Ready(Ok(n))
}
}
impl<R> NixReader<R>
where
R: AsyncReadExt + Unpin,
{
async fn force_fill(&mut self) -> io::Result<usize> {
let mut p = Pin::new(self);
let read = poll_fn(|cx| p.as_mut().poll_force_fill_buf(cx)).await?;
Ok(read)
}
}
impl<R> NixRead for NixReader<R>
where
R: AsyncReadExt + Send + Unpin,
{
type Error = io::Error;
fn version(&self) -> ProtocolVersion {
self.version
}
async fn try_read_number(&mut self) -> Result<Option<u64>, Self::Error> {
let mut buf = [0u8; 8];
let read = self.read_buf(&mut &mut buf[..]).await?;
if read == 0 {
return Ok(None);
}
if read < 8 {
self.read_exact(&mut buf[read..]).await?;
}
let num = Buf::get_u64_le(&mut &buf[..]);
Ok(Some(num))
}
async fn try_read_bytes_limited(
&mut self,
limit: RangeInclusive<usize>,
) -> Result<Option<Bytes>, Self::Error> {
assert!(
*limit.end() <= self.max_buf_size,
"The limit must be smaller than {}",
self.max_buf_size
);
match self.try_read_number().await? {
Some(raw_len) => {
// Check that length is in range and convert to usize
let len = raw_len
.try_into()
.ok()
.filter(|v| limit.contains(v))
.ok_or_else(|| Self::Error::invalid_data("bytes length out of range"))?;
// Calculate 64bit aligned length and convert to usize
let aligned: usize = raw_len
.checked_add(7)
.map(|v| v & !7)
.ok_or_else(|| Self::Error::invalid_data("bytes length out of range"))?
.try_into()
.map_err(Self::Error::invalid_data)?;
// Ensure that there is enough space in buffer for contents
if self.buf.len() + self.remaining_mut() < aligned {
self.buf.reserve(aligned - self.buf.len());
}
while self.buf.len() < aligned {
if self.force_fill().await? == 0 {
return Err(Self::Error::missing_data(
"unexpected end-of-file reading bytes",
));
}
}
let mut contents = self.buf.split_to(aligned);
let padding = aligned - len;
// Ensure padding is all zeros
if contents[len..] != EMPTY_BYTES[..padding] {
return Err(Self::Error::invalid_data("non-zero padding"));
}
contents.truncate(len);
Ok(Some(contents.freeze()))
}
None => Ok(None),
}
}
fn try_read_bytes(
&mut self,
) -> impl std::future::Future<Output = Result<Option<Bytes>, Self::Error>> + Send + '_ {
self.try_read_bytes_limited(0..=self.max_buf_size)
}
fn read_bytes(
&mut self,
) -> impl std::future::Future<Output = Result<Bytes, Self::Error>> + Send + '_ {
self.read_bytes_limited(0..=self.max_buf_size)
}
}
impl<R: AsyncRead> AsyncRead for NixReader<R> {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let rem = ready!(self.as_mut().poll_fill_buf(cx))?;
let amt = std::cmp::min(rem.len(), buf.remaining());
buf.put_slice(&rem[0..amt]);
self.consume(amt);
Poll::Ready(Ok(()))
}
}
impl<R: AsyncRead> AsyncBufRead for NixReader<R> {
fn poll_fill_buf(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
if self.as_ref().project_ref().buf.is_empty() {
ready!(self.as_mut().poll_force_fill_buf(cx))?;
}
let me = self.project();
Poll::Ready(Ok(&me.buf[..]))
}
fn consume(self: Pin<&mut Self>, amt: usize) {
let me = self.project();
me.buf.advance(amt)
}
}
#[cfg(test)]
mod test {
use std::time::Duration;
use hex_literal::hex;
use rstest::rstest;
use tokio_test::io::Builder;
use super::*;
use crate::wire::de::NixRead;
#[tokio::test]
async fn test_read_u64() {
let mock = Builder::new().read(&hex!("0100 0000 0000 0000")).build();
let mut reader = NixReader::new(mock);
assert_eq!(1, reader.read_number().await.unwrap());
assert_eq!(hex!(""), reader.buffer());
let mut buf = Vec::new();
reader.read_to_end(&mut buf).await.unwrap();
assert_eq!(hex!(""), &buf[..]);
}
#[tokio::test]
async fn test_read_u64_rest() {
let mock = Builder::new()
.read(&hex!("0100 0000 0000 0000 0123 4567 89AB CDEF"))
.build();
let mut reader = NixReader::new(mock);
assert_eq!(1, reader.read_number().await.unwrap());
assert_eq!(hex!("0123 4567 89AB CDEF"), reader.buffer());
let mut buf = Vec::new();
reader.read_to_end(&mut buf).await.unwrap();
assert_eq!(hex!("0123 4567 89AB CDEF"), &buf[..]);
}
#[tokio::test]
async fn test_read_u64_partial() {
let mock = Builder::new()
.read(&hex!("0100 0000"))
.wait(Duration::ZERO)
.read(&hex!("0000 0000 0123 4567 89AB CDEF"))
.wait(Duration::ZERO)
.read(&hex!("0100 0000"))
.build();
let mut reader = NixReader::new(mock);
assert_eq!(1, reader.read_number().await.unwrap());
assert_eq!(hex!("0123 4567 89AB CDEF"), reader.buffer());
let mut buf = Vec::new();
reader.read_to_end(&mut buf).await.unwrap();
assert_eq!(hex!("0123 4567 89AB CDEF 0100 0000"), &buf[..]);
}
#[tokio::test]
async fn test_read_u64_eof() {
let mock = Builder::new().build();
let mut reader = NixReader::new(mock);
assert_eq!(
io::ErrorKind::UnexpectedEof,
reader.read_number().await.unwrap_err().kind()
);
}
#[tokio::test]
async fn test_try_read_u64_none() {
let mock = Builder::new().build();
let mut reader = NixReader::new(mock);
assert_eq!(None, reader.try_read_number().await.unwrap());
}
#[tokio::test]
async fn test_try_read_u64_eof() {
let mock = Builder::new().read(&hex!("0100 0000 0000")).build();
let mut reader = NixReader::new(mock);
assert_eq!(
io::ErrorKind::UnexpectedEof,
reader.try_read_number().await.unwrap_err().kind()
);
}
#[tokio::test]
async fn test_try_read_u64_eof2() {
let mock = Builder::new()
.read(&hex!("0100"))
.wait(Duration::ZERO)
.read(&hex!("0000 0000"))
.build();
let mut reader = NixReader::new(mock);
assert_eq!(
io::ErrorKind::UnexpectedEof,
reader.try_read_number().await.unwrap_err().kind()
);
}
#[rstest]
#[case::empty(b"", &hex!("0000 0000 0000 0000"))]
#[case::one(b")", &hex!("0100 0000 0000 0000 2900 0000 0000 0000"))]
#[case::two(b"it", &hex!("0200 0000 0000 0000 6974 0000 0000 0000"))]
#[case::three(b"tea", &hex!("0300 0000 0000 0000 7465 6100 0000 0000"))]
#[case::four(b"were", &hex!("0400 0000 0000 0000 7765 7265 0000 0000"))]
#[case::five(b"where", &hex!("0500 0000 0000 0000 7768 6572 6500 0000"))]
#[case::six(b"unwrap", &hex!("0600 0000 0000 0000 756E 7772 6170 0000"))]
#[case::seven(b"where's", &hex!("0700 0000 0000 0000 7768 6572 6527 7300"))]
#[case::aligned(b"read_tea", &hex!("0800 0000 0000 0000 7265 6164 5F74 6561"))]
#[case::more_bytes(b"read_tess", &hex!("0900 0000 0000 0000 7265 6164 5F74 6573 7300 0000 0000 0000"))]
#[tokio::test]
async fn test_read_bytes(#[case] expected: &[u8], #[case] data: &[u8]) {
let mock = Builder::new().read(data).build();
let mut reader = NixReader::new(mock);
let actual = reader.read_bytes().await.unwrap();
assert_eq!(&actual[..], expected);
}
#[tokio::test]
async fn test_read_bytes_empty() {
let mock = Builder::new().build();
let mut reader = NixReader::new(mock);
assert_eq!(
io::ErrorKind::UnexpectedEof,
reader.read_bytes().await.unwrap_err().kind()
);
}
#[tokio::test]
async fn test_try_read_bytes_none() {
let mock = Builder::new().build();
let mut reader = NixReader::new(mock);
assert_eq!(None, reader.try_read_bytes().await.unwrap());
}
#[tokio::test]
async fn test_try_read_bytes_missing_data() {
let mock = Builder::new()
.read(&hex!("0500"))
.wait(Duration::ZERO)
.read(&hex!("0000 0000"))
.build();
let mut reader = NixReader::new(mock);
assert_eq!(
io::ErrorKind::UnexpectedEof,
reader.try_read_bytes().await.unwrap_err().kind()
);
}
#[tokio::test]
async fn test_try_read_bytes_missing_padding() {
let mock = Builder::new()
.read(&hex!("0200 0000 0000 0000"))
.wait(Duration::ZERO)
.read(&hex!("1234"))
.build();
let mut reader = NixReader::new(mock);
assert_eq!(
io::ErrorKind::UnexpectedEof,
reader.try_read_bytes().await.unwrap_err().kind()
);
}
#[tokio::test]
async fn test_read_bytes_bad_padding() {
let mock = Builder::new()
.read(&hex!("0200 0000 0000 0000"))
.wait(Duration::ZERO)
.read(&hex!("1234 0100 0000 0000"))
.build();
let mut reader = NixReader::new(mock);
assert_eq!(
io::ErrorKind::InvalidData,
reader.read_bytes().await.unwrap_err().kind()
);
}
#[tokio::test]
async fn test_read_bytes_limited_out_of_range() {
let mock = Builder::new().read(&hex!("FFFF 0000 0000 0000")).build();
let mut reader = NixReader::new(mock);
assert_eq!(
io::ErrorKind::InvalidData,
reader.read_bytes_limited(0..=50).await.unwrap_err().kind()
);
}
#[tokio::test]
async fn test_read_bytes_length_overflow() {
let mock = Builder::new().read(&hex!("F9FF FFFF FFFF FFFF")).build();
let mut reader = NixReader::builder()
.set_max_buf_size(usize::MAX)
.build(mock);
assert_eq!(
io::ErrorKind::InvalidData,
reader
.read_bytes_limited(0..=usize::MAX)
.await
.unwrap_err()
.kind()
);
}
// FUTUREWORK: Test this on supported hardware
#[tokio::test]
#[cfg(any(target_pointer_width = "16", target_pointer_width = "32"))]
async fn test_bytes_length_conversion_overflow() {
let len = (usize::MAX as u64) + 1;
let mock = Builder::new().read(&len.to_le_bytes()).build();
let mut reader = NixReader::new(mock);
assert_eq!(
std::io::ErrorKind::InvalidData,
reader.read_value::<usize>().await.unwrap_err().kind()
);
}
// FUTUREWORK: Test this on supported hardware
#[tokio::test]
#[cfg(any(target_pointer_width = "16", target_pointer_width = "32"))]
async fn test_bytes_aligned_length_conversion_overflow() {
let len = (usize::MAX - 6) as u64;
let mock = Builder::new().read(&len.to_le_bytes()).build();
let mut reader = NixReader::new(mock);
assert_eq!(
std::io::ErrorKind::InvalidData,
reader.read_value::<usize>().await.unwrap_err().kind()
);
}
#[tokio::test]
async fn test_buffer_resize() {
let mock = Builder::new()
.read(&hex!("0100"))
.read(&hex!("0000 0000 0000"))
.build();
let mut reader = NixReader::builder().set_reserved_buf_size(8).build(mock);
// buffer has no capacity initially
assert_eq!(0, reader.buffer_mut().capacity());
assert_eq!(2, reader.force_fill().await.unwrap());
// After first read buffer should have capacity we chose
assert_eq!(8, reader.buffer_mut().capacity());
// Because there was only 6 bytes remaining in buffer,
// which is enough to read the last 6 bytes, but we require
// capacity for 8 bytes, it doubled the capacity
assert_eq!(6, reader.force_fill().await.unwrap());
assert_eq!(16, reader.buffer_mut().capacity());
assert_eq!(1, reader.read_number().await.unwrap());
}
}
