use crate::decode::{SSP2_MAGIC, WIRE_VERSION};
use crate::error::{DecodeError, Result};
use crate::model::frame_type;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ScannedMessage {
pub message: Vec<u8>,
pub excess: usize,
}
#[derive(Default)]
pub struct MessageStreamScanner {
buffer: Vec<u8>,
offset: Option<usize>,
complete: bool,
}
impl MessageStreamScanner {
pub fn new() -> Self {
Self::default()
}
pub fn started(&self) -> bool {
!self.buffer.is_empty()
}
pub fn push(&mut self, chunk: &[u8]) -> Result<Option<ScannedMessage>> {
assert!(
!self.complete,
"MessageStreamScanner: message already complete"
);
self.buffer.extend_from_slice(chunk);
if self.offset.is_none() {
if self.buffer.len() < 8 {
return Ok(None);
}
self.check_header()?;
self.offset = Some(8);
}
let mut offset = self.offset.expect("offset set once header parsed");
loop {
if self.buffer.len() - offset < 5 {
self.offset = Some(offset);
return Ok(None);
}
let frame_type = self.buffer[offset];
let frame_length = u32::from_le_bytes([
self.buffer[offset + 1],
self.buffer[offset + 2],
self.buffer[offset + 3],
self.buffer[offset + 4],
]) as usize;
let frame_end = offset + 5 + frame_length;
if self.buffer.len() < frame_end {
self.offset = Some(offset);
return Ok(None);
}
offset = frame_end;
if frame_type == frame_type::END {
self.complete = true;
self.offset = Some(offset);
return Ok(Some(ScannedMessage {
message: self.buffer[..frame_end].to_vec(),
excess: self.buffer.len() - frame_end,
}));
}
}
}
fn check_header(&self) -> Result<()> {
let b = &self.buffer;
if b[0..4] != SSP2_MAGIC[..] {
return Err(DecodeError::invalid("bad envelope magic (expected SSP2)"));
}
let wire_version = u16::from_le_bytes([b[4], b[5]]);
if wire_version != WIRE_VERSION {
return Err(DecodeError::invalid(format!(
"unsupported wireVersion {wire_version}"
)));
}
if b[6] != 0x01 && b[6] != 0x02 {
return Err(DecodeError::invalid(format!(
"unknown msgKind byte 0x{:02x}",
b[6]
)));
}
if b[7] != 0x00 {
return Err(DecodeError::invalid(format!(
"envelope flags must be 0x00, got 0x{:02x}",
b[7]
)));
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::encode::encode_message;
use crate::model::{Frame, Message, MsgKind};
fn request_bytes() -> Vec<u8> {
let message = Message {
msg_kind: MsgKind::Request,
frames: vec![
Frame::ReqHeader {
client_id: "c1".to_owned(),
schema_version: 1,
},
Frame::PullHeader {
limit_commits: 0,
limit_snapshot_rows: 0,
max_snapshot_pages: 0,
accept: 0b0011,
},
],
};
encode_message(&message)
}
#[test]
fn whole_message_in_one_chunk_completes_with_zero_excess() {
let request = request_bytes();
let mut scanner = MessageStreamScanner::new();
let result = scanner.push(&request).unwrap().expect("complete");
assert_eq!(result.excess, 0);
assert_eq!(result.message, request);
}
#[test]
fn every_split_point_reassembles_byte_exactly() {
let request = request_bytes();
for split in 1..request.len() {
let mut scanner = MessageStreamScanner::new();
assert!(
scanner.push(&request[..split]).unwrap().is_none(),
"split {split}: first half must be incomplete"
);
let result = scanner
.push(&request[split..])
.unwrap()
.unwrap_or_else(|| panic!("split {split}: second half must complete"));
assert_eq!(result.excess, 0, "split {split}");
assert_eq!(result.message, request, "split {split}");
}
}
#[test]
fn one_byte_at_a_time_trickle_completes() {
let request = request_bytes();
let mut scanner = MessageStreamScanner::new();
let mut result = None;
for byte in &request {
result = scanner.push(&[*byte]).unwrap();
}
let result = result.expect("complete after last byte");
assert_eq!(result.excess, 0);
assert_eq!(result.message, request);
}
#[test]
fn bytes_past_end_are_reported_as_excess() {
let request = request_bytes();
let mut with_excess = request.clone();
with_excess.extend_from_slice(&[0xaa, 0xbb, 0xcc]);
let mut scanner = MessageStreamScanner::new();
let result = scanner.push(&with_excess).unwrap().expect("complete");
assert_eq!(result.excess, 3);
assert_eq!(result.message, request);
}
#[test]
fn two_message_stream_every_split_reassembles_first_exactly() {
let first = request_bytes();
let second = request_bytes();
let mut stream = first.clone();
stream.extend_from_slice(&second);
for split in 1..stream.len() {
let mut scanner = MessageStreamScanner::new();
let (result, buffered_at_completion) = match scanner.push(&stream[..split]).unwrap() {
Some(done) => (done, split),
None => {
let done = scanner
.push(&stream[split..])
.unwrap()
.unwrap_or_else(|| panic!("split {split}: first message must complete"));
(done, stream.len())
}
};
assert_eq!(result.message, first, "split {split}: first message bytes");
assert_eq!(
result.excess,
buffered_at_completion - first.len(),
"split {split}: excess equals second-message bytes buffered at END"
);
}
}
#[test]
fn bad_magic_is_a_decode_error_once_header_arrived() {
let mut request = request_bytes();
request[0] = 0x58;
let mut scanner = MessageStreamScanner::new();
assert!(scanner.push(&request[..4]).unwrap().is_none());
assert!(scanner.push(&request[4..]).is_err());
}
#[test]
fn non_zero_flags_and_unknown_msg_kind_are_decode_errors() {
for (index, value) in [(6usize, 0x03u8), (7usize, 0x01u8)] {
let mut request = request_bytes();
request[index] = value;
let mut scanner = MessageStreamScanner::new();
assert!(scanner.push(&request).is_err(), "index {index}");
}
}
#[test]
#[should_panic(expected = "already complete")]
fn push_after_completion_panics() {
let request = request_bytes();
let mut scanner = MessageStreamScanner::new();
scanner.push(&request).unwrap();
let _ = scanner.push(&[0]);
}
}