99 lines
3.0 KiB
Rust
99 lines
3.0 KiB
Rust
//! Length-prefixed frame codec: one envelope per big-endian 4-byte length.
|
|
|
|
use std::io::Read;
|
|
|
|
use crate::{Envelope, PROTOCOL_VERSION};
|
|
|
|
pub const MAX_FRAME: usize = 1_048_576;
|
|
|
|
#[derive(Debug)]
|
|
pub enum FrameError {
|
|
Closed,
|
|
Empty,
|
|
TooLarge(usize),
|
|
BadVersion(u32),
|
|
Json(serde_json::Error),
|
|
Io(std::io::Error),
|
|
}
|
|
|
|
impl std::fmt::Display for FrameError {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
FrameError::Closed => write!(f, "stream closed before a frame"),
|
|
FrameError::Empty => write!(f, "frame length was zero"),
|
|
FrameError::TooLarge(n) => write!(f, "frame of {n} bytes exceeds the maximum"),
|
|
FrameError::BadVersion(v) => write!(f, "unsupported protocol version {v}"),
|
|
FrameError::Json(e) => write!(f, "invalid JSON envelope: {e}"),
|
|
FrameError::Io(e) => write!(f, "frame I/O error: {e}"),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl std::error::Error for FrameError {
|
|
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
|
|
match self {
|
|
FrameError::Json(e) => Some(e),
|
|
FrameError::Io(e) => Some(e),
|
|
_ => None,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn write_frame<W: std::io::Write>(w: &mut W, env: &Envelope) -> Result<(), FrameError> {
|
|
let body = serde_json::to_vec(env).map_err(FrameError::Json)?;
|
|
let len = body.len();
|
|
if len > MAX_FRAME {
|
|
return Err(FrameError::TooLarge(len));
|
|
}
|
|
w.write_all(&(len as u32).to_be_bytes())
|
|
.map_err(FrameError::Io)?;
|
|
w.write_all(&body).map_err(FrameError::Io)?;
|
|
w.flush().map_err(FrameError::Io)
|
|
}
|
|
|
|
pub fn read_frame<R: std::io::Read>(r: &mut R) -> Result<Envelope, FrameError> {
|
|
let mut buf = [0u8; 4];
|
|
let mut filled = 0;
|
|
while filled < 4 {
|
|
let n = read_bytes(r, &mut buf[filled..]).map_err(FrameError::Io)?;
|
|
if n == 0 {
|
|
return Err(if filled == 0 {
|
|
FrameError::Closed
|
|
} else {
|
|
FrameError::Io(std::io::Error::from(std::io::ErrorKind::UnexpectedEof))
|
|
});
|
|
}
|
|
filled += n;
|
|
}
|
|
|
|
let len = u32::from_be_bytes(buf);
|
|
if len == 0 {
|
|
return Err(FrameError::Empty);
|
|
}
|
|
if len > MAX_FRAME as u32 {
|
|
return Err(FrameError::TooLarge(len as usize));
|
|
}
|
|
|
|
let len = len as usize;
|
|
let mut body = vec![0u8; len];
|
|
r.read_exact(&mut body).map_err(FrameError::Io)?;
|
|
|
|
let env = serde_json::from_slice::<Envelope>(&body).map_err(FrameError::Json)?;
|
|
if env.v != PROTOCOL_VERSION {
|
|
return Err(FrameError::BadVersion(env.v));
|
|
}
|
|
Ok(env)
|
|
}
|
|
|
|
/// Reads into `buf`, retrying on `Interrupted`. Returns `Ok(0)` on a clean end of stream.
|
|
fn read_bytes<R: Read>(r: &mut R, buf: &mut [u8]) -> std::io::Result<usize> {
|
|
loop {
|
|
match r.read(buf) {
|
|
Ok(0) => return Ok(0),
|
|
Ok(n) => return Ok(n),
|
|
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
|
|
Err(e) => return Err(e),
|
|
}
|
|
}
|
|
}
|