73 lines
2 KiB
Rust
73 lines
2 KiB
Rust
|
use futures::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
|
||
|
|
|
||
|
|
pub const VERSION: u8 = 1;
|
||
|
|
pub const MAX_FRAME: usize = 65535;
|
||
|
|
|
||
|
|
pub type Result<T> = std::result::Result<T, Error>;
|
||
|
|
|
||
|
|
#[derive(Debug)]
|
||
|
|
pub struct Error(pub String);
|
||
|
|
|
||
|
|
impl std::fmt::Display for Error {
|
||
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
|
|
self.0.fmt(f)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
impl std::error::Error for Error {}
|
||
|
|
|
||
|
|
pub fn encode(body: &[u8]) -> Result<Vec<u8>> {
|
||
|
|
if body.len() > MAX_FRAME {
|
||
|
|
return Err(Error("frame too large".into()));
|
||
|
|
}
|
||
|
|
let mut out = Vec::with_capacity(5 + body.len());
|
||
|
|
out.push(VERSION);
|
||
|
|
out.extend_from_slice(&(body.len() as u32).to_be_bytes());
|
||
|
|
out.extend_from_slice(body);
|
||
|
|
Ok(out)
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn decode(data: &[u8]) -> Result<Vec<u8>> {
|
||
|
|
if data.len() < 5 {
|
||
|
|
return Err(Error("short frame".into()));
|
||
|
|
}
|
||
|
|
if data[0] != VERSION {
|
||
|
|
return Err(Error("bad version".into()));
|
||
|
|
}
|
||
|
|
let len = u32::from_be_bytes(data[1..5].try_into().expect("4 bytes")) as usize;
|
||
|
|
if data.len() != 5 + len {
|
||
|
|
return Err(Error("length mismatch".into()));
|
||
|
|
}
|
||
|
|
if len > MAX_FRAME {
|
||
|
|
return Err(Error("frame too large".into()));
|
||
|
|
}
|
||
|
|
Ok(data[5..].to_vec())
|
||
|
|
}
|
||
|
|
|
||
|
|
pub async fn write_frame<W: AsyncWrite + Unpin>(w: &mut W, body: &[u8]) -> Result<()> {
|
||
|
|
let framed = encode(body)?;
|
||
|
|
w.write_all(&framed)
|
||
|
|
.await
|
||
|
|
.map_err(|e| Error(e.to_string()))?;
|
||
|
|
w.flush().await.map_err(|e| Error(e.to_string()))
|
||
|
|
}
|
||
|
|
|
||
|
|
pub async fn read_frame<R: AsyncRead + Unpin>(r: &mut R) -> Result<Vec<u8>> {
|
||
|
|
let mut hdr = [0u8; 5];
|
||
|
|
r.read_exact(&mut hdr)
|
||
|
|
.await
|
||
|
|
.map_err(|e| Error(e.to_string()))?;
|
||
|
|
if hdr[0] != VERSION {
|
||
|
|
return Err(Error("bad version".into()));
|
||
|
|
}
|
||
|
|
let len = u32::from_be_bytes(hdr[1..5].try_into().expect("4 bytes")) as usize;
|
||
|
|
if len > MAX_FRAME {
|
||
|
|
return Err(Error("frame too large".into()));
|
||
|
|
}
|
||
|
|
let mut body = vec![0u8; len];
|
||
|
|
r.read_exact(&mut body)
|
||
|
|
.await
|
||
|
|
.map_err(|e| Error(e.to_string()))?;
|
||
|
|
Ok(body)
|
||
|
|
}
|