use serde_json::Value;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader};
pub const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024;
#[derive(Debug)]
pub enum Inbound {
Frame(Value),
Violation(Violation),
}
#[derive(Debug, PartialEq, Eq)]
pub enum Violation {
TooLarge,
InvalidJson,
}
#[derive(Debug)]
pub struct FrameReader<R> {
reader: BufReader<R>,
max_frame: usize,
buf: Vec<u8>,
discarding: bool,
}
impl<R: AsyncRead + Unpin> FrameReader<R> {
pub fn new(reader: R, max_frame: usize) -> Self {
debug_assert!(max_frame > 0);
Self {
reader: BufReader::new(reader),
max_frame,
buf: Vec::new(),
discarding: false,
}
}
pub fn into_inner(self) -> BufReader<R> {
self.reader
}
pub async fn next(&mut self) -> std::io::Result<Option<Inbound>> {
let mut byte = [0u8; 1];
loop {
let n = self.reader.read(&mut byte).await?;
if n == 0 {
if std::mem::take(&mut self.discarding) {
return Ok(Some(Inbound::Violation(Violation::TooLarge)));
}
return Ok(None);
}
if byte[0] == b'\n' {
if std::mem::take(&mut self.discarding) {
return Ok(Some(Inbound::Violation(Violation::TooLarge)));
}
let line = std::mem::take(&mut self.buf);
return Ok(Some(match serde_json::from_slice::<Value>(&line) {
Ok(v) => Inbound::Frame(v),
Err(_) => Inbound::Violation(Violation::InvalidJson),
}));
}
if self.discarding {
continue;
}
if self.buf.len() >= self.max_frame {
self.discarding = true;
self.buf.clear();
continue;
}
self.buf.push(byte[0]);
}
}
}
pub async fn write_frame<W: AsyncWrite + Unpin>(w: &mut W, frame: &Value) -> std::io::Result<()> {
let mut line = serde_json::to_vec(frame)?;
line.push(b'\n');
w.write_all(&line).await?;
w.flush().await
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn reader_over(bytes: &[u8], cap: usize) -> FrameReader<&[u8]> {
FrameReader::new(bytes, cap)
}
#[tokio::test]
async fn valid_frames_roundtrip() {
let mut buf = Vec::new();
write_frame(&mut buf, &json!({"jsonrpc":"2.0","id":1,"method":"ping"}))
.await
.unwrap();
write_frame(&mut buf, &json!({"jsonrpc":"2.0","id":2,"method":"pong"}))
.await
.unwrap();
let mut r = reader_over(&buf, 1024);
assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 1));
assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 2));
assert!(r.next().await.unwrap().is_none()); }
#[tokio::test]
async fn oversized_frame_is_discarded_and_reported_and_stream_continues() {
let big = format!("{{\"pad\":\"{}\"}}\n", "x".repeat(100));
let mut bytes = big.into_bytes();
bytes.extend_from_slice(b"{\"jsonrpc\":\"2.0\",\"id\":7,\"method\":\"ok\"}\n");
let mut r = reader_over(&bytes, 64);
assert!(matches!(
r.next().await.unwrap().unwrap(),
Inbound::Violation(Violation::TooLarge)
));
assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 7));
}
#[tokio::test]
async fn invalid_json_is_a_violation() {
let mut r = reader_over(b"not json at all\n", 1024);
assert!(matches!(
r.next().await.unwrap().unwrap(),
Inbound::Violation(Violation::InvalidJson)
));
}
#[tokio::test]
async fn oversized_frame_truncated_by_eof_is_still_reported() {
let bytes = format!("{{\"pad\":\"{}\"}}", "x".repeat(100)).into_bytes(); let mut r = reader_over(&bytes, 64);
assert!(matches!(
r.next().await.unwrap().unwrap(),
Inbound::Violation(Violation::TooLarge)
));
assert!(r.next().await.unwrap().is_none());
}
#[tokio::test]
async fn back_to_back_oversized_frames_then_valid_frame() {
let big = format!("{{\"pad\":\"{}\"}}\n", "x".repeat(100));
let mut bytes = big.clone().into_bytes();
bytes.extend_from_slice(big.as_bytes());
bytes.extend_from_slice(b"{\"jsonrpc\":\"2.0\",\"id\":9,\"method\":\"ok\"}\n");
let mut r = reader_over(&bytes, 64);
assert!(matches!(
r.next().await.unwrap().unwrap(),
Inbound::Violation(Violation::TooLarge)
));
assert!(matches!(
r.next().await.unwrap().unwrap(),
Inbound::Violation(Violation::TooLarge)
));
assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 9));
}
#[tokio::test]
async fn frame_of_exactly_max_frame_bytes_passes_one_more_is_too_large() {
let payload = "x".repeat(62);
let exact = format!("\"{payload}\"\n");
let mut r = reader_over(exact.as_bytes(), 64);
assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v == payload));
let over = format!("\"{}\"\n", "x".repeat(63)); let mut r = reader_over(over.as_bytes(), 64);
assert!(matches!(
r.next().await.unwrap().unwrap(),
Inbound::Violation(Violation::TooLarge)
));
}
#[tokio::test]
async fn into_inner_preserves_pipelined_read_ahead() {
let bytes = b"{\"id\":1}\n{\"id\":2}\n";
let mut r = reader_over(bytes, 1024);
assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 1));
let boxed: Box<dyn AsyncRead + Unpin + Send> = Box::new(r.into_inner());
let mut r2 = FrameReader::new(boxed, 1024);
assert!(matches!(r2.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 2));
assert!(r2.next().await.unwrap().is_none());
}
}