use core::future::Future;
use tokio::io::{AsyncRead, AsyncReadExt};
use tokio::sync::Mutex;
use crate::application::jsonrpc::{FrameSource, SourceError};
use crate::domain::{FrameDecoder, MaxFrameSize};
use super::StdioConfig;
#[derive(Debug)]
struct Inner<R> {
stream: R,
decoder: FrameDecoder,
buf: Vec<u8>,
ended: bool,
failed: bool,
}
#[derive(Debug)]
pub struct StdioSource<R> {
inner: Mutex<Inner<R>>,
config: StdioConfig,
}
impl<R> StdioSource<R> {
#[must_use]
pub fn new(stream: R, config: StdioConfig) -> Self {
let buf = vec![0; config.read_chunk().get()];
Self {
inner: Mutex::new(Inner {
stream,
decoder: FrameDecoder::new(config.limit()),
buf,
ended: false,
failed: false,
}),
config,
}
}
#[must_use]
pub const fn config(&self) -> StdioConfig {
self.config
}
#[must_use]
pub const fn limit(&self) -> MaxFrameSize {
self.config.limit()
}
}
impl<R> StdioSource<R>
where
R: AsyncRead + Unpin,
{
#[expect(
clippy::significant_drop_tightening,
reason = "the guard is the mutual exclusion for the entire read loop, and every path out of that loop is still holding it"
)]
pub async fn next_frame(&self) -> Result<Option<String>, SourceError> {
let mut guard = self.inner.lock().await;
let Inner {
stream,
decoder,
buf,
ended,
failed,
} = &mut *guard;
if *failed {
return Err(SourceError::Transport);
}
loop {
if let Some(frame) = decoder.next_frame()? {
return Ok(Some(frame.to_owned()));
}
if *ended {
return Ok(None);
}
let Ok(read) = stream.read(buf.as_mut_slice()).await else {
*failed = true;
return Err(SourceError::Transport);
};
if read == 0 {
*ended = true;
return Ok(None);
}
let Some(bytes) = buf.get(..read) else {
*failed = true;
return Err(SourceError::Transport);
};
decoder.push(bytes)?;
}
}
}
impl<R> FrameSource for StdioSource<R>
where
R: AsyncRead + Unpin + Send,
{
fn next_frame(&self) -> impl Future<Output = Result<Option<String>, SourceError>> + Send {
Self::next_frame(self)
}
}
#[cfg(test)]
mod tests {
use core::pin::Pin;
use core::sync::atomic::{AtomicUsize, Ordering};
use core::task::{Context, Poll};
use std::io;
use std::sync::Arc;
use assert_matches::assert_matches;
use tokio::io::{AsyncRead, ReadBuf};
use super::{StdioConfig, StdioSource};
use crate::application::jsonrpc::{DynFrameSource, SourceError};
use crate::domain::{FrameDecoder, FramingError, MaxFrameSize, MaxFrameSizeError};
#[derive(Debug)]
#[expect(dead_code, reason = "the payloads are read by the derived Debug")]
enum Failure {
Source(SourceError),
Limit(MaxFrameSizeError),
}
impl From<SourceError> for Failure {
fn from(error: SourceError) -> Self {
Self::Source(error)
}
}
impl From<MaxFrameSizeError> for Failure {
fn from(error: MaxFrameSizeError) -> Self {
Self::Limit(error)
}
}
struct Trickle {
data: Vec<u8>,
served: usize,
per_read: usize,
fail_after: Option<usize>,
reads: Arc<AtomicUsize>,
}
impl Trickle {
fn new(data: &[u8], per_read: usize) -> (Self, Arc<AtomicUsize>) {
let reads = Arc::new(AtomicUsize::new(0));
(
Self {
data: data.to_vec(),
served: 0,
per_read,
fail_after: None,
reads: Arc::clone(&reads),
},
reads,
)
}
fn failing(data: &[u8], per_read: usize) -> Self {
Self {
fail_after: Some(0),
..Self::new(data, per_read).0
}
}
}
impl AsyncRead for Trickle {
fn poll_read(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let this = self.get_mut();
this.reads.fetch_add(1, Ordering::Relaxed);
if this.fail_after.is_some_and(|limit| this.served >= limit) {
return Poll::Ready(Err(io::Error::other("injected read failure")));
}
let remaining = this.data.len().saturating_sub(this.served);
if remaining == 0 {
return Poll::Ready(Ok(()));
}
let end = this
.served
.saturating_add(this.per_read)
.min(this.data.len());
let Some(chunk) = this.data.get(this.served..end) else {
return Poll::Ready(Ok(()));
};
buf.put_slice(chunk);
this.served = end;
Poll::Ready(Ok(()))
}
}
fn trickling(
wire: &[u8],
per_read: usize,
limit: u32,
) -> Result<(StdioSource<Trickle>, Arc<AtomicUsize>), Failure> {
let (stream, reads) = Trickle::new(wire, per_read);
let config = StdioConfig::new(MaxFrameSize::new(limit)?).with_read_chunk(per_read);
Ok((StdioSource::new(stream, config), reads))
}
#[tokio::test]
async fn a_frame_split_across_several_reads_reassembles() -> Result<(), Failure> {
let document = r#"{"jsonrpc":"2.0","id":1,"result":{}}"#;
let mut wire = document.as_bytes().to_vec();
wire.push(b'\n');
let (source, reads) = trickling(&wire, 4, 1024)?;
assert_eq!(source.next_frame().await?, Some(document.to_owned()));
assert!(
reads.load(Ordering::Relaxed) > 1,
"the frame must have spanned more than one read for this test to mean anything"
);
Ok(())
}
#[tokio::test]
async fn a_frame_buffered_before_the_close_is_still_delivered() -> Result<(), Failure> {
let (source, _) = trickling(b"one\ntwo\n", 8, 1024)?;
assert_eq!(source.next_frame().await?, Some("one".to_owned()));
assert_eq!(source.next_frame().await?, Some("two".to_owned()));
assert_eq!(source.next_frame().await?, None);
Ok(())
}
#[tokio::test]
async fn a_clean_end_of_stream_is_a_close_and_stays_one() -> Result<(), Failure> {
let (source, _) = trickling(b"{\"jsonrpc\":\"2.0\"}\n", 64, 1024)?;
assert!(source.next_frame().await?.is_some());
assert_eq!(source.next_frame().await?, None);
assert_eq!(source.next_frame().await?, None);
Ok(())
}
#[tokio::test]
async fn a_read_failure_is_a_transport_error() -> Result<(), Failure> {
let config = StdioConfig::new(MaxFrameSize::new(1024)?);
let source = StdioSource::new(Trickle::failing(b"ignored\n", 8), config);
assert_matches!(source.next_frame().await, Err(SourceError::Transport));
Ok(())
}
#[tokio::test]
async fn a_failed_read_stays_failed_and_never_becomes_a_clean_close() -> Result<(), Failure> {
let config = StdioConfig::new(MaxFrameSize::new(1024)?);
let source = StdioSource::new(Trickle::failing(b"ignored\n", 8), config);
for _ in 0..3 {
assert_matches!(source.next_frame().await, Err(SourceError::Transport));
}
Ok(())
}
#[tokio::test]
async fn an_oversized_inbound_frame_is_refused() -> Result<(), Failure> {
let (source, _) = trickling(&[b'x'; 200], 16, 64)?;
assert_matches!(
source.next_frame().await,
Err(SourceError::Framing(FramingError::FrameTooLarge { .. }))
);
Ok(())
}
#[tokio::test]
async fn an_unterminated_trailing_line_is_not_a_frame() -> Result<(), Failure> {
let (source, _) = trickling(b"{\"jsonrpc\":\"2.0\"}", 64, 1024)?;
assert_eq!(source.next_frame().await?, None);
Ok(())
}
#[tokio::test]
async fn non_utf8_bytes_are_refused() -> Result<(), Failure> {
let (source, _) = trickling(&[0xFF, 0xFE, b'\n'], 8, 1024)?;
assert_matches!(
source.next_frame().await,
Err(SourceError::Framing(FramingError::NotUtf8 { .. }))
);
Ok(())
}
#[tokio::test]
async fn the_source_is_usable_behind_a_trait_object() -> Result<(), Failure> {
let (source, _) = trickling(b"behind-a-box\n", 16, 1024)?;
let boxed: Arc<dyn DynFrameSource> = Arc::new(source);
assert_eq!(boxed.next_frame().await?, Some("behind-a-box".to_owned()));
Ok(())
}
#[test]
fn the_framing_is_the_domain_decoder_not_a_local_copy() {
let mut decoder =
FrameDecoder::new(MaxFrameSize::new(1024).unwrap_or(MaxFrameSize::DEFAULT));
assert!(decoder.is_empty());
assert_matches!(decoder.next_frame(), Ok(None));
}
}