use std::{
io::Result,
pin::Pin,
task::{ready, Poll},
};
use pin_project_lite::pin_project;
use tokio::io::{AsyncRead, ReadBuf};
/// State machine for [`NixFramedReader`].
///
/// As the reader progresses it linearly cycles through the states.
#[derive(Debug)]
enum NixFramedReaderState {
/// The reader always starts in this state.
///
/// Before the payload, the client first sends its size.
/// The size is a u64 which is 8 bytes long, while it's likely that we will receive
/// the whole u64 in one read, it's possible that it will arrive in smaller chunks.
/// So in this state we read up to 8 bytes and transition to
/// [`NixFramedReaderState::ReadingPayload`] when done if the read size is not zero,
/// otherwise we reset filled to 0, and read the next size value.
ReadingSize { buf: [u8; 8], filled: usize },
/// This is where we read the actual payload that is sent to us.
///
/// Once we've read the expected number of bytes, we go back to the
/// [`NixFramedReaderState::ReadingSize`] state.
ReadingPayload {
/// Represents the remaining number of bytes we expect to read based on the value
/// read in the previous state.
remaining: u64,
},
}
pin_project! {
/// Implements Nix's Framed reader protocol for protocol versions >= 1.23.
///
/// See serialization.md#framed and [`NixFramedReaderState`] for details.
pub struct NixFramedReader<R> {
#[pin]
reader: R,
state: NixFramedReaderState,
}
}
impl<R> NixFramedReader<R> {
pub fn new(reader: R) -> Self {
Self {
reader,
state: NixFramedReaderState::ReadingSize {
buf: [0; 8],
filled: 0,
},
}
}
}
impl<R: AsyncRead> AsyncRead for NixFramedReader<R> {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
read_buf: &mut ReadBuf<'_>,
) -> Poll<Result<()>> {
let mut this = self.as_mut().project();
match this.state {
NixFramedReaderState::ReadingSize { buf, filled } => {
if *filled < buf.len() {
let mut size_buf = ReadBuf::new(buf);
size_buf.advance(*filled);
ready!(this.reader.poll_read(cx, &mut size_buf))?;
let bytes_read = size_buf.filled().len() - *filled;
if bytes_read == 0 {
// oef
return Poll::Ready(Ok(()));
}
*filled += bytes_read;
// Schedule ourselves to run again.
return self.poll_read(cx, read_buf);
}
let size = u64::from_le_bytes(*buf);
if size == 0 {
// eof
*filled = 0;
return Poll::Ready(Ok(()));
}
*this.state = NixFramedReaderState::ReadingPayload { remaining: size };
self.poll_read(cx, read_buf)
}
NixFramedReaderState::ReadingPayload { remaining } => {
// Make sure we never try to read more than usize which is 4 bytes on 32-bit platforms.
let safe_remaining = if *remaining <= usize::MAX as u64 {
*remaining as usize
} else {
usize::MAX
};
if safe_remaining > 0 {
// The buffer is no larger than the amount of data that we expect.
// Otherwise we will trim the buffer below and come back here.
if read_buf.remaining() <= safe_remaining {
let filled_before = read_buf.filled().len();
ready!(this.reader.as_mut().poll_read(cx, read_buf))?;
let bytes_read = read_buf.filled().len() - filled_before;
*remaining -= bytes_read as u64;
if *remaining == 0 {
*this.state = NixFramedReaderState::ReadingSize {
buf: [0; 8],
filled: 0,
};
}
return Poll::Ready(Ok(()));
}
// Don't read more than remaining + pad bytes, it avoids unnecessary allocations and makes
// internal bookkeeping simpler.
let mut smaller_buf = read_buf.take(safe_remaining);
ready!(self.as_mut().poll_read(cx, &mut smaller_buf))?;
let bytes_read = smaller_buf.filled().len();
// SAFETY: we just read this number of bytes into read_buf's backing slice above.
unsafe { read_buf.assume_init(bytes_read) };
read_buf.advance(bytes_read);
return Poll::Ready(Ok(()));
}
*this.state = NixFramedReaderState::ReadingSize {
buf: [0; 8],
filled: 0,
};
self.poll_read(cx, read_buf)
}
}
}
}
#[cfg(test)]
mod nix_framed_tests {
use std::time::Duration;
use tokio::io::AsyncReadExt;
use tokio_test::io::Builder;
use crate::nix_daemon::framing::NixFramedReader;
#[tokio::test]
async fn read_hello_world_in_two_frames() {
let mut mock = Builder::new()
// The client sends len
.read(&5u64.to_le_bytes())
// Immediately followed by the bytes
.read("hello".as_bytes())
.wait(Duration::ZERO)
// Send more data separately
.read(&6u64.to_le_bytes())
.read(" world".as_bytes())
.build();
let mut reader = NixFramedReader::new(&mut mock);
let mut result = String::new();
reader
.read_to_string(&mut result)
.await
.expect("Could not read into result");
assert_eq!("hello world", result);
}
#[tokio::test]
async fn read_hello_world_in_two_frames_followed_by_zero_sized_frame() {
let mut mock = Builder::new()
// The client sends len
.read(&5u64.to_le_bytes())
// Immediately followed by the bytes
.read("hello".as_bytes())
.wait(Duration::ZERO)
// Send more data separately
.read(&6u64.to_le_bytes())
.read(" world".as_bytes())
.read(&0u64.to_le_bytes())
.build();
let mut reader = NixFramedReader::new(&mut mock);
let mut result = String::new();
reader
.read_to_string(&mut result)
.await
.expect("Could not read into result");
assert_eq!("hello world", result);
}
}
