use crate::codec::{read_bytes, read_u32_le, read_u64_le};
use crate::error::WireError;
use bytes::{BufMut, BytesMut};
use std::borrow::Cow;
use std::num::NonZeroU32;
pub const STATUS_OK: u32 = 0;
#[derive(Debug)]
pub struct RequestFrame<'a> {
pub code: u32,
pub request_id: u64,
pub payload: &'a [u8],
}
impl<'a> RequestFrame<'a> {
pub const HEADER_SIZE: usize = 16;
pub fn payload_length(frame_length: u32) -> Result<u32, WireError> {
frame_length
.checked_sub(12)
.ok_or(WireError::Validation(Cow::Borrowed(
"request frame length must be at least 12 (code + request_id)",
)))
}
#[must_use]
pub const fn from_parts(code: u32, request_id: u64, payload: &'a [u8]) -> Self {
Self {
code,
request_id,
payload,
}
}
pub fn decode(buf: &'a [u8]) -> Result<(Self, usize), WireError> {
let frame_length = read_u32_le(buf, 0)?;
let payload_len = Self::payload_length(frame_length)? as usize;
let code = read_u32_le(buf, 4)?;
let request_id = read_u64_le(buf, 8)?;
let payload = read_bytes(buf, Self::HEADER_SIZE, payload_len)?;
let total = Self::HEADER_SIZE + payload_len;
Ok((
Self {
code,
request_id,
payload,
},
total,
))
}
pub fn encode(
code: u32,
request_id: u64,
payload: &[u8],
out: &mut BytesMut,
) -> Result<(), WireError> {
let length = payload
.len()
.checked_add(12)
.and_then(|n| u32::try_from(n).ok())
.ok_or(WireError::PayloadTooLarge {
size: payload.len(),
max: u32::MAX as usize - 12,
})?;
out.reserve(Self::HEADER_SIZE + payload.len());
out.put_u32_le(length);
out.put_u32_le(code);
out.put_u64_le(request_id);
out.put_slice(payload);
Ok(())
}
#[must_use]
pub const fn encoded_size(payload_len: usize) -> Option<usize> {
Self::HEADER_SIZE.checked_add(payload_len)
}
}
#[derive(Debug)]
pub struct ResponseFrame<'a> {
pub status: u32,
pub request_id: u64,
pub payload: &'a [u8],
}
impl<'a> ResponseFrame<'a> {
pub const HEADER_SIZE: usize = 16;
pub fn decode(buf: &'a [u8]) -> Result<(Self, usize), WireError> {
let status = read_u32_le(buf, 0)?;
let length = read_u32_le(buf, 4)? as usize;
if length < 8 {
return Err(WireError::Validation(Cow::Borrowed(
"response frame length must be at least 8 (request_id)",
)));
}
let request_id = read_u64_le(buf, 8)?;
let payload_len = length - 8;
let payload = read_bytes(buf, Self::HEADER_SIZE, payload_len)?;
let total = Self::HEADER_SIZE + payload_len;
Ok((
Self {
status,
request_id,
payload,
},
total,
))
}
pub fn encode_ok(request_id: u64, payload: &[u8], out: &mut BytesMut) -> Result<(), WireError> {
let length = payload
.len()
.checked_add(8)
.and_then(|n| u32::try_from(n).ok())
.ok_or(WireError::PayloadTooLarge {
size: payload.len(),
max: u32::MAX as usize - 8,
})?;
out.reserve(Self::HEADER_SIZE + payload.len());
out.put_u32_le(STATUS_OK);
out.put_u32_le(length);
out.put_u64_le(request_id);
out.put_slice(payload);
Ok(())
}
pub fn encode_error(status: NonZeroU32, request_id: u64, out: &mut BytesMut) {
out.reserve(Self::HEADER_SIZE);
out.put_u32_le(status.get());
out.put_u32_le(8); out.put_u64_le(request_id);
}
#[must_use]
pub const fn is_ok(&self) -> bool {
self.status == STATUS_OK
}
#[must_use]
pub const fn encoded_size(payload_len: usize) -> Option<usize> {
Self::HEADER_SIZE.checked_add(payload_len)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn request2_roundtrip() {
let payload = b"hello world";
let mut buf = BytesMut::with_capacity(RequestFrame::encoded_size(payload.len()).unwrap());
RequestFrame::encode(42, 7, payload, &mut buf).unwrap();
let (frame, consumed) = RequestFrame::decode(&buf).unwrap();
assert_eq!(consumed, buf.len());
assert_eq!(frame.code, 42);
assert_eq!(frame.request_id, 7);
assert_eq!(frame.payload, payload);
}
#[test]
fn request2_empty_payload() {
let mut buf = BytesMut::with_capacity(RequestFrame::HEADER_SIZE);
RequestFrame::encode(1, 99, &[], &mut buf).unwrap();
let (frame, consumed) = RequestFrame::decode(&buf).unwrap();
assert_eq!(consumed, 16);
assert_eq!(frame.code, 1);
assert_eq!(frame.request_id, 99);
assert!(frame.payload.is_empty());
}
#[test]
fn request2_length_field_includes_code_and_request_id() {
let payload = b"test";
let mut buf = BytesMut::new();
RequestFrame::encode(99, 1, payload, &mut buf).unwrap();
let length = u32::from_le_bytes(buf[0..4].try_into().unwrap());
assert_eq!(length, 4 + 8 + 4); }
#[test]
fn request2_truncated_header() {
let buf = [0u8; 15]; assert!(RequestFrame::decode(&buf).is_err());
}
#[test]
fn request2_truncated_payload() {
let mut buf = BytesMut::new();
buf.put_u32_le(112); buf.put_u32_le(1); buf.put_u64_le(1); buf.put_slice(&[0u8; 50]); assert!(RequestFrame::decode(&buf).is_err());
}
#[test]
fn request2_length_too_small() {
let mut buf = BytesMut::new();
buf.put_u32_le(11); buf.put_u32_le(1);
buf.put_u64_le(1);
assert!(RequestFrame::decode(&buf).is_err());
}
#[test]
fn request2_payload_length_valid() {
assert_eq!(RequestFrame::payload_length(12).unwrap(), 0);
assert_eq!(RequestFrame::payload_length(112).unwrap(), 100);
}
#[test]
fn request2_payload_length_too_small() {
assert!(RequestFrame::payload_length(0).is_err());
assert!(RequestFrame::payload_length(11).is_err());
}
#[test]
fn request2_encoded_size() {
assert_eq!(RequestFrame::encoded_size(0), Some(16));
assert_eq!(RequestFrame::encoded_size(100), Some(116));
assert_eq!(RequestFrame::encoded_size(usize::MAX), None);
}
#[test]
fn request2_from_parts() {
let payload = b"data";
let frame = RequestFrame::from_parts(5, 42, payload);
assert_eq!(frame.code, 5);
assert_eq!(frame.request_id, 42);
assert_eq!(frame.payload, payload);
}
#[test]
fn response2_ok_roundtrip() {
let payload = b"response data";
let mut buf = BytesMut::with_capacity(ResponseFrame::encoded_size(payload.len()).unwrap());
ResponseFrame::encode_ok(7, payload, &mut buf).unwrap();
let (frame, consumed) = ResponseFrame::decode(&buf).unwrap();
assert_eq!(consumed, buf.len());
assert!(frame.is_ok());
assert_eq!(frame.request_id, 7);
assert_eq!(frame.payload, payload);
}
#[test]
fn response2_ok_empty_payload() {
let mut buf = BytesMut::new();
ResponseFrame::encode_ok(42, &[], &mut buf).unwrap();
let (frame, consumed) = ResponseFrame::decode(&buf).unwrap();
assert_eq!(consumed, 16);
assert!(frame.is_ok());
assert_eq!(frame.request_id, 42);
assert!(frame.payload.is_empty());
}
#[test]
fn response2_error_roundtrip() {
let mut buf = BytesMut::new();
ResponseFrame::encode_error(NonZeroU32::new(1001).unwrap(), 55, &mut buf);
let (frame, consumed) = ResponseFrame::decode(&buf).unwrap();
assert_eq!(consumed, 16);
assert!(!frame.is_ok());
assert_eq!(frame.status, 1001);
assert_eq!(frame.request_id, 55);
assert!(frame.payload.is_empty());
}
#[test]
fn response2_truncated_header() {
let buf = [0u8; 15];
assert!(ResponseFrame::decode(&buf).is_err());
}
#[test]
fn response2_length_too_small() {
let mut buf = BytesMut::new();
buf.put_u32_le(0); buf.put_u32_le(7); buf.put_u64_le(1); assert!(ResponseFrame::decode(&buf).is_err());
}
#[test]
fn response2_encoded_size() {
assert_eq!(ResponseFrame::encoded_size(0), Some(16));
assert_eq!(ResponseFrame::encoded_size(256), Some(272));
assert_eq!(ResponseFrame::encoded_size(usize::MAX), None);
}
}