use alloc::vec::Vec;
use byteorder::{BigEndian, ByteOrder};
use tracing::trace;
#[derive(Debug)]
pub struct TcpBuffer {
buf: Vec<u8>,
}
impl core::fmt::Display for TcpBuffer {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("TcpBuffer")
.field("buf", &alloc::format!("{} bytes", self.buf.len()))
.finish()
}
}
impl TcpBuffer {
pub fn new() -> Self {
Vec::new().into()
}
pub fn push_data(&mut self, data: &[u8]) {
self.buf.extend(data);
}
pub fn pull_data(&mut self) -> Option<Vec<u8>> {
if self.buf.len() < 2 {
trace!(
"running buffer is currently too small ({} bytes) to provide data",
self.buf.len()
);
return None;
}
let data_length = BigEndian::read_u16(&self.buf[..2]) as usize;
if self.buf.len() < data_length {
trace!(
"not enough data, buf length {} data specifies length {}",
self.buf.len(),
data_length
);
return None;
}
let bytes = self.take(data_length);
trace!("return {} bytes", data_length);
Some(bytes)
}
fn take(&mut self, data_length: usize) -> Vec<u8> {
let offset = data_length + 2;
if offset > self.buf.len() {
return Vec::new();
}
let mut data = self.buf.split_off(offset);
core::mem::swap(&mut data, &mut self.buf);
data[2..].to_vec()
}
}
impl Default for TcpBuffer {
fn default() -> Self {
Self::new()
}
}
impl From<Vec<u8>> for TcpBuffer {
fn from(value: Vec<u8>) -> Self {
Self { buf: value }
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn tcp_buffer_split_recv() {
let _log = crate::tests::test_init_log();
let mut tcp_buffer = TcpBuffer::default();
let mut len = [0; 2];
let data = [0, 1, 2, 4, 3];
BigEndian::write_u16(&mut len, data.len() as u16);
tcp_buffer.push_data(&len);
assert!(tcp_buffer.pull_data().is_none());
tcp_buffer.push_data(&data);
assert_eq!(tcp_buffer.pull_data().unwrap(), &data);
}
}