mirror of https://github.com/ospab/ostp.git
116 lines
3.1 KiB
Rust
116 lines
3.1 KiB
Rust
use bytes::{BufMut, Bytes, BytesMut};
|
|
|
|
use crate::protocol::ProtocolError;
|
|
|
|
const FRAME_HEADER_LEN: usize = 12;
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
#[repr(u8)]
|
|
pub enum FrameKind {
|
|
Handshake = 1,
|
|
Data = 2,
|
|
Close = 3,
|
|
KeepAlive = 4,
|
|
Nack = 5,
|
|
Ack = 6,
|
|
}
|
|
|
|
impl TryFrom<u8> for FrameKind {
|
|
type Error = ProtocolError;
|
|
|
|
fn try_from(value: u8) -> Result<Self, Self::Error> {
|
|
match value {
|
|
1 => Ok(Self::Handshake),
|
|
2 => Ok(Self::Data),
|
|
3 => Ok(Self::Close),
|
|
4 => Ok(Self::KeepAlive),
|
|
5 => Ok(Self::Nack),
|
|
6 => Ok(Self::Ack),
|
|
_ => Err(ProtocolError::Framing("unknown frame kind".to_string())),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub struct FrameHeader {
|
|
pub version: u8,
|
|
pub kind: FrameKind,
|
|
pub stream_id: u16,
|
|
pub payload_len: u32,
|
|
pub pad_len: u16,
|
|
}
|
|
|
|
impl FrameHeader {
|
|
pub fn encode(&self, out: &mut BytesMut) {
|
|
out.put_u8(self.version);
|
|
out.put_u8(self.kind as u8);
|
|
out.put_u16(0); // 2 reserved bytes
|
|
out.put_u16(self.stream_id);
|
|
out.put_u32(self.payload_len);
|
|
out.put_u16(self.pad_len);
|
|
}
|
|
|
|
pub fn decode(buf: &[u8]) -> Result<Self, ProtocolError> {
|
|
if buf.len() < FRAME_HEADER_LEN {
|
|
return Err(ProtocolError::Framing("truncated frame header".to_string()));
|
|
}
|
|
|
|
let version = buf[0];
|
|
let kind = FrameKind::try_from(buf[1])?;
|
|
// buf[2] and buf[3] are reserved
|
|
let stream_id = u16::from_be_bytes([buf[4], buf[5]]);
|
|
let payload_len = u32::from_be_bytes([buf[6], buf[7], buf[8], buf[9]]);
|
|
let pad_len = u16::from_be_bytes([buf[10], buf[11]]);
|
|
|
|
Ok(Self {
|
|
version,
|
|
kind,
|
|
stream_id,
|
|
payload_len,
|
|
pad_len,
|
|
})
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct FramedPacket {
|
|
pub header: FrameHeader,
|
|
pub payload: Bytes,
|
|
pub padding: Bytes,
|
|
}
|
|
|
|
impl FramedPacket {
|
|
pub fn encode(&self) -> Bytes {
|
|
let total = FRAME_HEADER_LEN + self.payload.len() + self.padding.len();
|
|
let mut out = BytesMut::with_capacity(total);
|
|
self.header.encode(&mut out);
|
|
out.extend_from_slice(&self.payload);
|
|
out.extend_from_slice(&self.padding);
|
|
out.freeze()
|
|
}
|
|
|
|
pub fn decode_zero_copy(buf: Bytes) -> Result<Self, ProtocolError> {
|
|
if buf.len() < FRAME_HEADER_LEN {
|
|
return Err(ProtocolError::Framing("frame too short".to_string()));
|
|
}
|
|
|
|
let header = FrameHeader::decode(&buf[..FRAME_HEADER_LEN])?;
|
|
let payload_len = header.payload_len as usize;
|
|
let pad_len = header.pad_len as usize;
|
|
|
|
let expected = FRAME_HEADER_LEN + payload_len + pad_len;
|
|
if buf.len() < expected {
|
|
return Err(ProtocolError::Framing("frame body truncated".to_string()));
|
|
}
|
|
|
|
let payload = buf.slice(FRAME_HEADER_LEN..FRAME_HEADER_LEN + payload_len);
|
|
let padding = buf.slice(FRAME_HEADER_LEN + payload_len..expected);
|
|
|
|
Ok(Self {
|
|
header,
|
|
payload,
|
|
padding,
|
|
})
|
|
}
|
|
}
|