use bytes::{Buf, BufMut, Bytes, BytesMut};
pub const STREAM_HEADER_SIZE: usize = 8;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct StreamFlags(u8);
impl StreamFlags {
const KNOWN_MASK: u8 = Self::FIN.0 | Self::RST.0 | Self::ACK.0 | Self::SYN.0;
pub const NONE: StreamFlags = StreamFlags(0);
pub const FIN: StreamFlags = StreamFlags(1);
pub const RST: StreamFlags = StreamFlags(2);
pub const ACK: StreamFlags = StreamFlags(4);
pub const SYN: StreamFlags = StreamFlags(8);
pub fn from_byte(b: u8) -> Self {
StreamFlags(b)
}
pub fn as_byte(self) -> u8 {
self.0
}
pub fn is_empty(self) -> bool {
self.0 == 0
}
pub fn has_unknown_bits(self) -> bool {
self.0 & !Self::KNOWN_MASK != 0
}
pub fn is_fin(self) -> bool {
self.0 & Self::FIN.0 != 0
}
pub fn is_rst(self) -> bool {
self.0 & Self::RST.0 != 0
}
pub fn is_ack(self) -> bool {
self.0 & Self::ACK.0 != 0
}
pub fn is_syn(self) -> bool {
self.0 & Self::SYN.0 != 0
}
pub fn with(self, other: StreamFlags) -> StreamFlags {
StreamFlags(self.0 | other.0)
}
}
impl std::ops::BitOr for StreamFlags {
type Output = Self;
fn bitor(self, rhs: Self) -> Self::Output {
StreamFlags(self.0 | rhs.0)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct StreamHeader {
pub stream_id: u16,
pub flags: StreamFlags,
pub reserved: u8,
pub length: u32,
}
impl StreamHeader {
pub const CONTROL_STREAM: u16 = 0;
pub const MAX_STREAM_ID: u16 = 65535;
pub fn new(stream_id: u16, flags: StreamFlags, length: u32) -> Self {
Self {
stream_id,
flags,
reserved: 0,
length,
}
}
pub fn data(stream_id: u16, length: u32) -> Self {
Self::new(stream_id, StreamFlags::NONE, length)
}
pub fn syn(stream_id: u16) -> Self {
Self::new(stream_id, StreamFlags::SYN, 0)
}
pub fn fin(stream_id: u16) -> Self {
Self::new(stream_id, StreamFlags::FIN, 0)
}
pub fn rst(stream_id: u16) -> Self {
Self::new(stream_id, StreamFlags::RST, 0)
}
pub fn ack(stream_id: u16) -> Self {
Self::new(stream_id, StreamFlags::ACK, 0)
}
pub fn is_control(&self) -> bool {
self.stream_id == Self::CONTROL_STREAM
}
pub fn encode(&self, dst: &mut BytesMut) {
dst.put_u16(self.stream_id);
dst.put_u8(self.flags.as_byte());
dst.put_u8(self.reserved);
dst.put_u32(self.length);
}
pub fn to_bytes(&self) -> Bytes {
let mut buf = BytesMut::with_capacity(STREAM_HEADER_SIZE);
self.encode(&mut buf);
buf.freeze()
}
pub fn decode(src: &[u8]) -> Option<Self> {
if src.len() < STREAM_HEADER_SIZE {
return None;
}
Some(Self {
stream_id: u16::from_be_bytes([src[0], src[1]]),
flags: StreamFlags::from_byte(src[2]),
reserved: src[3],
length: u32::from_be_bytes([src[4], src[5], src[6], src[7]]),
})
}
pub fn decode_from(src: &mut BytesMut) -> Option<Self> {
if src.len() < STREAM_HEADER_SIZE {
return None;
}
let header = Self::decode(src)?;
src.advance(STREAM_HEADER_SIZE);
Some(header)
}
pub fn frame_size(&self) -> usize {
STREAM_HEADER_SIZE + self.length as usize
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_stream_flags() {
let flags = StreamFlags::FIN | StreamFlags::ACK;
assert!(flags.is_fin());
assert!(flags.is_ack());
assert!(!flags.is_rst());
assert!(!flags.is_syn());
}
#[test]
fn test_stream_header_encode_decode() {
let header = StreamHeader::new(42, StreamFlags::FIN, 1234);
let mut buf = BytesMut::new();
header.encode(&mut buf);
assert_eq!(buf.len(), STREAM_HEADER_SIZE);
let decoded = StreamHeader::decode(&buf).unwrap();
assert_eq!(decoded.stream_id, 42);
assert!(decoded.flags.is_fin());
assert_eq!(decoded.length, 1234);
}
#[test]
fn test_stream_header_helpers() {
let syn = StreamHeader::syn(1);
assert!(syn.flags.is_syn());
assert_eq!(syn.stream_id, 1);
let fin = StreamHeader::fin(2);
assert!(fin.flags.is_fin());
assert_eq!(fin.stream_id, 2);
let rst = StreamHeader::rst(3);
assert!(rst.flags.is_rst());
assert_eq!(rst.stream_id, 3);
}
#[test]
fn test_control_stream() {
let header = StreamHeader::data(StreamHeader::CONTROL_STREAM, 100);
assert!(header.is_control());
let header = StreamHeader::data(1, 100);
assert!(!header.is_control());
}
#[test]
fn test_frame_size() {
let header = StreamHeader::data(1, 100);
assert_eq!(header.frame_size(), STREAM_HEADER_SIZE + 100);
}
}