use tokio::io::{AsyncWrite, AsyncWriteExt};
use tokio::sync::Mutex;
use crate::application::session::{FrameSink, SinkError};
use crate::domain::{FramingError, MaxFrameSize, encode_frame};
use super::{StdioConfig, StdioError};
#[derive(Debug)]
pub struct StdioSink<W> {
stream: Mutex<W>,
limit: MaxFrameSize,
}
impl<W> StdioSink<W> {
#[must_use]
pub fn new(stream: W, config: StdioConfig) -> Self {
Self {
stream: Mutex::new(stream),
limit: config.limit(),
}
}
#[must_use]
pub const fn limit(&self) -> MaxFrameSize {
self.limit
}
}
impl<W> StdioSink<W>
where
W: AsyncWrite + Unpin,
{
pub async fn send(&self, document: &str) -> Result<(), StdioError> {
let mut framed = Vec::with_capacity(document.len() + 1);
encode_frame(document, &mut framed, self.limit).map_err(stdio_error)?;
let mut stream = self.stream.lock().await;
stream
.write_all(&framed)
.await
.map_err(|_| StdioError::Transport)?;
stream.flush().await.map_err(|_| StdioError::Transport)
}
}
const fn stdio_error(error: FramingError) -> StdioError {
match error {
FramingError::FrameTooLarge { limit, at_least } => StdioError::ReplyTooLarge {
limit,
len: at_least,
},
FramingError::NotUtf8 { .. } => StdioError::Transport,
}
}
impl<W> FrameSink for StdioSink<W>
where
W: AsyncWrite + Unpin + Send,
{
async fn send(&self, frame: &str) -> Result<(), SinkError> {
Self::send(self, frame).await.map_err(|_| SinkError)
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use assert_matches::assert_matches;
use super::super::{StdioConfig, testing::Capture};
use super::{StdioError, StdioSink, stdio_error};
use crate::application::session::{DynFrameSink, FrameSink, SinkError};
use crate::domain::{FramingError, MaxFrameSize, MaxFrameSizeError};
#[derive(Debug)]
#[expect(dead_code, reason = "the payloads are read by the derived Debug")]
enum Failure {
Stdio(StdioError),
Limit(MaxFrameSizeError),
Sink(SinkError),
}
impl From<StdioError> for Failure {
fn from(error: StdioError) -> Self {
Self::Stdio(error)
}
}
impl From<MaxFrameSizeError> for Failure {
fn from(error: MaxFrameSizeError) -> Self {
Self::Limit(error)
}
}
impl From<SinkError> for Failure {
fn from(error: SinkError) -> Self {
Self::Sink(error)
}
}
fn capture_sink(limit: u32) -> Result<(StdioSink<Capture>, Capture), Failure> {
let capture = Capture::shared();
let config = StdioConfig::new(MaxFrameSize::new(limit)?);
Ok((StdioSink::new(capture.clone(), config), capture))
}
fn ceiling(bytes: u32) -> MaxFrameSize {
MaxFrameSize::new(bytes).unwrap_or(MaxFrameSize::DEFAULT)
}
#[tokio::test]
async fn a_reply_is_written_as_exactly_one_line() -> Result<(), Failure> {
let (sink, capture) = capture_sink(1024)?;
let document = r#"{"jsonrpc":"2.0","id":1,"result":null}"#;
sink.send(document).await?;
assert_eq!(capture.bytes(), format!("{document}\n").into_bytes());
Ok(())
}
#[tokio::test]
async fn an_oversized_reply_is_typed_and_writes_nothing() -> Result<(), Failure> {
let (sink, capture) = capture_sink(16)?;
let result = sink
.send("a reply well over the sixteen byte ceiling")
.await;
assert_matches!(
result,
Err(StdioError::ReplyTooLarge { limit, len }) => {
assert_eq!(limit, ceiling(16));
assert!(len > 16, "the reported length is the document's, not the limit's");
}
);
assert!(
capture.bytes().is_empty(),
"an oversized reply must not put a single byte on the wire"
);
Ok(())
}
#[tokio::test]
async fn a_reply_of_exactly_the_ceiling_is_written() -> Result<(), Failure> {
let (sink, capture) = capture_sink(16)?;
let document = "0123456789abcdef";
sink.send(document).await?;
assert_eq!(
capture.bytes().len(),
17,
"sixteen document bytes plus one terminator"
);
Ok(())
}
#[tokio::test]
async fn the_port_trait_turns_an_oversized_reply_into_a_sink_error() -> Result<(), Failure> {
let (sink, capture) = capture_sink(16)?;
let result = FrameSink::send(&sink, "a reply well over the ceiling").await;
assert_matches!(result, Err(SinkError));
assert!(capture.bytes().is_empty());
Ok(())
}
#[tokio::test]
async fn concurrent_replies_do_not_interleave() -> Result<(), Failure> {
let (sink, capture) = capture_sink(1024)?;
let first = Arc::new(sink);
let second = Arc::clone(&first);
let (left, right) = tokio::join!(
async move { first.send("{\"id\":1}\nembedded").await },
async move { second.send("{\"id\":2}").await }
);
left?;
right?;
let written = capture.contents();
for document in ["{\"id\":1}\nembedded", "{\"id\":2}"] {
assert!(
written.contains(&format!("{document}\n")),
"each reply must land whole; the stream was {written:?}"
);
}
Ok(())
}
#[tokio::test]
async fn the_sink_is_usable_behind_a_trait_object() -> Result<(), Failure> {
let (sink, capture) = capture_sink(1024)?;
let boxed: Arc<dyn DynFrameSink> = Arc::new(sink);
DynFrameSink::send(boxed.as_ref(), "boxed").await?;
assert_eq!(capture.contents(), "boxed\n");
Ok(())
}
#[test]
fn a_not_utf8_framing_error_maps_rather_than_panics() {
assert_eq!(
stdio_error(FramingError::NotUtf8 { valid_up_to: 3 }),
StdioError::Transport
);
assert_matches!(
stdio_error(FramingError::FrameTooLarge {
limit: ceiling(64),
at_least: 65,
}),
StdioError::ReplyTooLarge { len: 65, .. }
);
}
}