use bytes::{Buf, BufMut, Bytes, BytesMut};
use std::io;
pub const FRAME_HEADER_SIZE: usize = 7;
#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Command {
Waste = 0,
Syn = 1,
Psh = 2,
Fin = 3,
Settings = 4,
Alert = 5,
UpdatePaddingScheme = 6,
SynAck = 7,
HeartRequest = 8,
HeartResponse = 9,
ServerSettings = 10,
}
impl TryFrom<u8> for Command {
type Error = io::Error;
fn try_from(value: u8) -> Result<Self, Self::Error> {
match value {
0 => Ok(Command::Waste),
1 => Ok(Command::Syn),
2 => Ok(Command::Psh),
3 => Ok(Command::Fin),
4 => Ok(Command::Settings),
5 => Ok(Command::Alert),
6 => Ok(Command::UpdatePaddingScheme),
7 => Ok(Command::SynAck),
8 => Ok(Command::HeartRequest),
9 => Ok(Command::HeartResponse),
10 => Ok(Command::ServerSettings),
_ => Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("unknown command: {}", value),
)),
}
}
}
#[derive(Debug, Clone)]
pub struct Frame {
pub cmd: Command,
pub stream_id: u32,
pub data: Bytes,
}
impl Frame {
pub fn control(cmd: Command, stream_id: u32) -> Self {
Self {
cmd,
stream_id,
data: Bytes::new(),
}
}
pub fn with_data(cmd: Command, stream_id: u32, data: Bytes) -> Self {
Self {
cmd,
stream_id,
data,
}
}
pub fn data(stream_id: u32, data: Bytes) -> Self {
Self::with_data(Command::Psh, stream_id, data)
}
#[cfg(test)]
pub fn encode(&self) -> BytesMut {
let mut buf = BytesMut::with_capacity(FRAME_HEADER_SIZE + self.data.len());
self.encode_into(&mut buf);
buf
}
#[inline]
pub fn encode_into(&self, buf: &mut BytesMut) {
buf.reserve(FRAME_HEADER_SIZE + self.data.len());
buf.put_u8(self.cmd as u8);
buf.put_u32(self.stream_id);
buf.put_u16(self.data.len() as u16);
if !self.data.is_empty() {
buf.extend_from_slice(&self.data);
}
}
pub fn decode_header(header: &[u8; FRAME_HEADER_SIZE]) -> io::Result<(Command, u32, u16)> {
let cmd = Command::try_from(header[0])?;
let stream_id = u32::from_be_bytes([header[1], header[2], header[3], header[4]]);
let length = u16::from_be_bytes([header[5], header[6]]);
Ok((cmd, stream_id, length))
}
}
#[derive(Debug, Clone, Default)]
pub struct StringMap(std::collections::HashMap<String, String>);
impl StringMap {
pub fn new() -> Self {
Self(std::collections::HashMap::new())
}
pub fn insert(&mut self, key: impl Into<String>, value: impl Into<String>) {
self.0.insert(key.into(), value.into());
}
pub fn get(&self, key: &str) -> Option<&String> {
self.0.get(key)
}
pub fn from_bytes(data: &[u8]) -> Self {
let mut map = Self::new();
let text = String::from_utf8_lossy(data);
for line in text.lines() {
if let Some((key, value)) = line.split_once('=') {
map.insert(key.to_string(), value.to_string());
}
}
map
}
pub fn to_bytes(&self) -> Vec<u8> {
let lines: Vec<String> = self.0.iter().map(|(k, v)| format!("{}={}", k, v)).collect();
lines.join("\n").into_bytes()
}
}
pub struct FrameCodec;
impl FrameCodec {
pub fn decode(buf: &mut BytesMut) -> io::Result<Option<Frame>> {
if buf.len() < FRAME_HEADER_SIZE {
return Ok(None);
}
let header: [u8; FRAME_HEADER_SIZE] = buf[..FRAME_HEADER_SIZE].try_into().unwrap();
let (cmd, stream_id, data_len) = Frame::decode_header(&header)?;
let total_len = FRAME_HEADER_SIZE + data_len as usize;
if buf.len() < total_len {
return Ok(None);
}
buf.advance(FRAME_HEADER_SIZE);
let data = if data_len > 0 {
buf.split_to(data_len as usize).freeze()
} else {
Bytes::new()
};
Ok(Some(Frame {
cmd,
stream_id,
data,
}))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_command_conversion() {
assert_eq!(Command::try_from(0).unwrap(), Command::Waste);
assert_eq!(Command::try_from(1).unwrap(), Command::Syn);
assert_eq!(Command::try_from(2).unwrap(), Command::Psh);
assert!(Command::try_from(255).is_err());
}
#[test]
fn test_frame_encode_decode() {
let frame = Frame::data(123, Bytes::from("hello"));
let encoded = frame.encode();
assert_eq!(encoded.len(), FRAME_HEADER_SIZE + 5);
let mut buf = encoded;
let decoded = FrameCodec::decode(&mut buf).unwrap().unwrap();
assert_eq!(decoded.cmd, Command::Psh);
assert_eq!(decoded.stream_id, 123);
assert_eq!(decoded.data.as_ref(), b"hello");
}
#[test]
fn test_frame_control() {
let frame = Frame::control(Command::Syn, 42);
assert_eq!(frame.data.len(), 0);
assert!(frame.data.is_empty());
}
#[test]
fn test_string_map() {
let mut map = StringMap::new();
map.insert("v", "2");
map.insert("client", "test");
let bytes = map.to_bytes();
let parsed = StringMap::from_bytes(&bytes);
assert_eq!(parsed.get("v"), Some(&"2".to_string()));
assert_eq!(parsed.get("client"), Some(&"test".to_string()));
}
}