diff --git a/protosocket-connection/src/connection.rs b/protosocket-connection/src/connection.rs index d13c649..ff17a9d 100644 --- a/protosocket-connection/src/connection.rs +++ b/protosocket-connection/src/connection.rs @@ -215,7 +215,7 @@ impl< Err(e) => match e { DeserializeError::IncompleteBuffer { next_message_size } => { if self.max_buffer_length < next_message_size { - log::error!( + log::info!( "tried to receive message that is too long. Resetting connection - max: {}, requested: {}", self.max_buffer_length, next_message_size @@ -474,25 +474,230 @@ impl< #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))] fn poll_receive(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<()> { loop { - match self.poll_read_inbound(context) { + break match self.poll_read_inbound(context) { ReadBufferState::Pending => { log::debug!("consumed all that I can from the read stream for now {self}"); - return Poll::Pending; + Poll::Pending } ReadBufferState::MoreToRead => { log::debug!("more to read"); - self.read_inbound_messages_and_react(); - continue; + match self.read_inbound_messages_and_react() { + ReadBufferState::Pending => { + // we want to park on inbound pending + continue; + } + ReadBufferState::MoreToRead => { + // partial message, still want to park on inbound pending + continue; + } + ReadBufferState::Error(e) => { + log::warn!("error while reading inbound message from stream: {e:?}"); + Poll::Ready(()) + } + ReadBufferState::Disconnected => { + log::info!("read connection closed while reading inbound message"); + Poll::Ready(()) + } + } } ReadBufferState::Disconnected => { log::info!("read connection closed"); - return Poll::Ready(()); + Poll::Ready(()) } ReadBufferState::Error(e) => { log::warn!("error while reading from tcp stream: {e:?}"); - return Poll::Ready(()); + Poll::Ready(()) } + }; + } + } +} + +#[cfg(test)] +mod test { + use std::{ + collections::VecDeque, + future::Future, + io::Cursor, + pin::Pin, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, + task::{Context, Poll, Waker}, + }; + + use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; + + use crate::{ + Codec, Connection, Decoder, DeserializeError, Encoder, MessageReactor, ReactorStatus, + }; + + const BUFFER_ALLOCATION_INCREMENT: usize = 8; + + /// Yields the scripted chunks, then parks forever. + #[derive(Debug)] + struct ScriptedStream { + inbound: VecDeque>, + } + + impl AsyncRead for ScriptedStream { + fn poll_read( + mut self: Pin<&mut Self>, + _context: &mut Context<'_>, + buffer: &mut ReadBuf<'_>, + ) -> Poll> { + match self.inbound.pop_front() { + Some(chunk) => { + assert!( + chunk.len() <= buffer.remaining(), + "scripted chunk must fit in the receive buffer" + ); + buffer.put_slice(&chunk); + Poll::Ready(Ok(())) + } + None => Poll::Pending, } } } + + impl AsyncWrite for ScriptedStream { + fn poll_write( + self: Pin<&mut Self>, + _context: &mut Context<'_>, + buffer: &[u8], + ) -> Poll> { + Poll::Ready(Ok(buffer.len())) + } + + fn poll_flush( + self: Pin<&mut Self>, + _context: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown( + self: Pin<&mut Self>, + _context: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + } + + /// Never produces a message: it always announces that it needs `next_message_size` bytes. + #[derive(Debug)] + struct IncompleteCodec { + next_message_size: usize, + decode_calls: Arc, + } + + impl Encoder for IncompleteCodec { + type Message = (); + type Serialized = Cursor>; + + fn encode(&mut self, _message: Self::Message) -> Self::Serialized { + Cursor::new(Vec::new()) + } + } + + impl Decoder for IncompleteCodec { + type Message = (); + + fn decode( + &mut self, + _buffer: impl bytes::Buf, + ) -> std::result::Result<(usize, Self::Message), DeserializeError> { + let calls = 1 + self.decode_calls.fetch_add(1, Ordering::Relaxed); + assert!( + calls < 1000, + "connection spun on the receive buffer: {calls} decode attempts in one poll" + ); + Err(DeserializeError::IncompleteBuffer { + next_message_size: self.next_message_size, + }) + } + } + + impl Codec for IncompleteCodec {} + + #[derive(Debug)] + struct NoopReactor; + + impl MessageReactor for NoopReactor { + type Codec = IncompleteCodec; + type Inbound = (); + type Outbound = (); + type LogicalOutbound = (); + + fn on_inbound_message(&mut self, _message: Self::Inbound) -> ReactorStatus { + ReactorStatus::Continue + } + + fn on_outbound_message(&mut self, message: Self::LogicalOutbound) -> Self::Outbound { + message + } + } + + #[allow(clippy::type_complexity)] + fn connection( + inbound: Vec, + max_buffer_length: usize, + next_message_size: usize, + decode_calls: Arc, + ) -> ( + spillway::Sender<()>, + Connection, + ) { + let (sender, receiver) = spillway::channel(); + let connection = Connection::new( + ScriptedStream { + inbound: VecDeque::from([inbound]), + }, + IncompleteCodec { + next_message_size, + decode_calls, + }, + max_buffer_length, + BUFFER_ALLOCATION_INCREMENT, + 1, + receiver, + NoopReactor, + ); + (sender, connection) + } + + /// A peer announcing a message larger than max_buffer_length used to wedge the poll loop: + /// the decode pass's Disconnected state was dropped, and the now-full receive buffer + /// returned MoreToRead forever. + #[test] + fn oversized_next_message_disconnects() { + let decode_calls = Arc::new(AtomicUsize::new(0)); + // fills the receive buffer exactly, so no further read can make room + let (_sender, mut connection) = connection( + vec![0; BUFFER_ALLOCATION_INCREMENT], + BUFFER_ALLOCATION_INCREMENT, + BUFFER_ALLOCATION_INCREMENT + 1, + decode_calls.clone(), + ); + + assert_eq!( + Poll::Ready(()), + Pin::new(&mut connection).poll(&mut Context::from_waker(Waker::noop())) + ); + assert_eq!(1, decode_calls.load(Ordering::Relaxed)); + } + + /// A message that has not arrived in full parks the connection rather than closing it. + #[test] + fn incomplete_next_message_parks() { + let decode_calls = Arc::new(AtomicUsize::new(0)); + let (_sender, mut connection) = connection(vec![0; 2], 64, 4, decode_calls.clone()); + + assert_eq!( + Poll::Pending, + Pin::new(&mut connection).poll(&mut Context::from_waker(Waker::noop())) + ); + assert_eq!(1, decode_calls.load(Ordering::Relaxed)); + } } diff --git a/protosocket-connection/src/error.rs b/protosocket-connection/src/error.rs index 8b8c8be..cfc47a5 100644 --- a/protosocket-connection/src/error.rs +++ b/protosocket-connection/src/error.rs @@ -4,9 +4,12 @@ pub enum DeserializeError { /// Buffer will be retained and you will be called again later with more bytes #[error("Need more bytes to decode the next message")] IncompleteBuffer { - /// This is a hint to the connection for how many more bytes should be read. + /// Total bytes the next message occupies on the wire, inclusive of framing. /// You may be called again before you get another buffer with at least this /// many bytes. + /// + /// The connection disconnects when this exceeds its max buffer length. + /// Undercounting wedges a connection on a message it can never buffer. next_message_size: usize, }, /// Buffer will be discarded diff --git a/protosocket-prost/src/decoder.rs b/protosocket-prost/src/decoder.rs index c6bdc6e..f62adf7 100644 --- a/protosocket-prost/src/decoder.rs +++ b/protosocket-prost/src/decoder.rs @@ -19,10 +19,10 @@ where ) -> std::result::Result<(usize, Self::Message), DeserializeError> { match prost::decode_length_delimiter(buffer.chunk()) { Ok(message_length) => { - if buffer.remaining() < message_length + prost::length_delimiter_len(message_length) - { + let framed_length = message_length + prost::length_delimiter_len(message_length); + if buffer.remaining() < framed_length { return Err(DeserializeError::IncompleteBuffer { - next_message_size: message_length, + next_message_size: framed_length, }); } }