use bytes::{Buf, Bytes, BytesMut};
use crate::errors::CoreError;
use crate::transport::session_transport::FramePhase;
pub(crate) const CT_CHANGE_CIPHER_SPEC: u8 = 0x14;
pub(crate) const CT_HANDSHAKE: u8 = 0x16;
pub(crate) const CT_APPLICATION_DATA: u8 = 0x17;
pub(crate) const VER_TLS10: u16 = 0x0301;
pub(crate) const VER_TLS12: u16 = 0x0303;
pub(crate) const TLS_RECORD_HEADER_LEN: usize = 5;
pub(crate) const MAX_RECORD_FRAGMENT: usize = 16384;
pub(crate) const MAX_RECORD_FRAGMENT_WIRE: usize = MAX_RECORD_FRAGMENT + 256;
const INNER_CHUNK_HEADER_LEN: usize = 2;
const HANDSHAKE_FRAME_CAP: usize = 64 * 1024;
const STEADY_STATE_FRAME_CAP: usize = 4 * 1024 * 1024;
const STREAM_SLACK: usize = 256 * 1024;
const MAX_CONSECUTIVE_EMPTY: u32 = 64;
const MAX_STREAM_CHUNK: usize = MAX_RECORD_FRAGMENT - INNER_CHUNK_HEADER_LEN;
pub(crate) fn encode_record(
content_type: u8,
version: u16,
fragment: &[u8],
out: &mut Vec<u8>,
) -> Result<(), CoreError> {
if fragment.len() > MAX_RECORD_FRAGMENT_WIRE {
return Err(CoreError::NetworkError(format!(
"TLS record fragment too large: {} > {}",
fragment.len(),
MAX_RECORD_FRAGMENT_WIRE
)));
}
out.push(content_type);
out.extend_from_slice(&version.to_be_bytes());
out.extend_from_slice(&(fragment.len() as u16).to_be_bytes());
out.extend_from_slice(fragment);
Ok(())
}
pub(crate) fn frame_message(message: &[u8]) -> Result<Vec<u8>, CoreError> {
let msg_len = u32::try_from(message.len()).map_err(|_| {
CoreError::NetworkError(format!("message too large to frame: {}", message.len()))
})?;
let mut stream = Vec::with_capacity(4 + message.len());
stream.extend_from_slice(&msg_len.to_be_bytes());
stream.extend_from_slice(message);
let mut out = Vec::new();
for chunk in stream.chunks(MAX_STREAM_CHUNK) {
let mut fragment = Vec::with_capacity(INNER_CHUNK_HEADER_LEN + chunk.len());
fragment.extend_from_slice(&(chunk.len() as u16).to_be_bytes());
fragment.extend_from_slice(chunk);
encode_record(CT_APPLICATION_DATA, VER_TLS12, &fragment, &mut out)?;
}
Ok(out)
}
pub(crate) struct RecordDeframer {
wire: BytesMut,
stream: BytesMut,
cap: usize,
consecutive_empty: u32,
}
impl Default for RecordDeframer {
fn default() -> Self {
Self::new()
}
}
impl RecordDeframer {
pub(crate) fn new() -> Self {
Self {
wire: BytesMut::new(),
stream: BytesMut::new(),
cap: HANDSHAKE_FRAME_CAP,
consecutive_empty: 0,
}
}
pub(crate) fn set_phase(&mut self, phase: FramePhase) {
self.cap = match phase {
FramePhase::Handshake => HANDSHAKE_FRAME_CAP,
FramePhase::Established => STEADY_STATE_FRAME_CAP,
};
}
pub(crate) fn push(&mut self, data: &[u8]) {
self.wire.extend_from_slice(data);
}
fn drain_records(&mut self) -> Result<(), CoreError> {
loop {
if self.wire.len() < TLS_RECORD_HEADER_LEN {
return Ok(()); }
let content_type = self.wire[0];
let frag_len = u16::from_be_bytes([self.wire[3], self.wire[4]]) as usize;
if frag_len > MAX_RECORD_FRAGMENT_WIRE {
return Err(CoreError::NetworkError(format!(
"oversized TLS record from peer: {} > {}",
frag_len, MAX_RECORD_FRAGMENT_WIRE
)));
}
if content_type != CT_APPLICATION_DATA {
return Err(CoreError::NetworkError(format!(
"unexpected TLS content_type in data phase: 0x{content_type:02x}"
)));
}
if self.wire.len() < TLS_RECORD_HEADER_LEN + frag_len {
return Ok(()); }
self.wire.advance(TLS_RECORD_HEADER_LEN);
let fragment = self.wire.split_to(frag_len);
if frag_len == 0 {
self.note_empty()?;
continue;
}
if frag_len < INNER_CHUNK_HEADER_LEN {
return Err(CoreError::NetworkError(
"TLS record fragment too small for inner framing".into(),
));
}
let chunk_len = u16::from_be_bytes([fragment[0], fragment[1]]) as usize;
if chunk_len > frag_len - INNER_CHUNK_HEADER_LEN {
return Err(CoreError::NetworkError(format!(
"inner chunk_len {} exceeds fragment body {}",
chunk_len,
frag_len - INNER_CHUNK_HEADER_LEN
)));
}
if chunk_len == 0 {
self.note_empty()?;
continue;
}
self.consecutive_empty = 0;
self.stream.extend_from_slice(
&fragment[INNER_CHUNK_HEADER_LEN..INNER_CHUNK_HEADER_LEN + chunk_len],
);
if self.stream.len() > self.cap.saturating_add(STREAM_SLACK) {
return Err(CoreError::NetworkError(format!(
"inner reassembly buffer exceeded cap: {} > {} + {}",
self.stream.len(),
self.cap,
STREAM_SLACK
)));
}
}
}
fn note_empty(&mut self) -> Result<(), CoreError> {
self.consecutive_empty += 1;
if self.consecutive_empty > MAX_CONSECUTIVE_EMPTY {
return Err(CoreError::NetworkError(
"too many consecutive empty TLS records (livelock defense)".into(),
));
}
Ok(())
}
pub(crate) fn next_message(&mut self) -> Result<Option<Bytes>, CoreError> {
self.drain_records()?;
if self.stream.len() < 4 {
return Ok(None);
}
let msg_len = u32::from_be_bytes([
self.stream[0],
self.stream[1],
self.stream[2],
self.stream[3],
]) as usize;
if msg_len > self.cap {
return Err(CoreError::NetworkError(format!(
"oversized inner message from peer: {} > {}",
msg_len, self.cap
)));
}
if self.stream.len() < 4 + msg_len {
return Ok(None);
}
self.stream.advance(4);
Ok(Some(self.stream.split_to(msg_len).freeze()))
}
}
#[cfg(test)]
mod tests {
use super::*;
const CT_HANDSHAKE: u8 = 0x16;
fn padded_record(chunk: &[u8], pad: usize) -> Vec<u8> {
let mut fragment = Vec::new();
fragment.extend_from_slice(&(chunk.len() as u16).to_be_bytes());
fragment.extend_from_slice(chunk);
fragment.extend(std::iter::repeat(0u8).take(pad));
let mut out = Vec::new();
encode_record(CT_APPLICATION_DATA, VER_TLS12, &fragment, &mut out).expect("encode");
out
}
#[test]
fn round_trip_single_message() {
let msg = b"hello phantom over tls";
let wire = frame_message(msg).expect("frame");
let mut d = RecordDeframer::new();
d.push(&wire);
let got = d.next_message().expect("parse").expect("one message");
assert_eq!(&got[..], msg);
assert!(
d.next_message().expect("parse").is_none(),
"no second message"
);
}
#[test]
fn round_trip_multiple_messages_concatenated() {
let mut wire = Vec::new();
wire.extend_from_slice(&frame_message(b"first").expect("f1"));
wire.extend_from_slice(&frame_message(b"second").expect("f2"));
wire.extend_from_slice(&frame_message(b"third").expect("f3"));
let mut d = RecordDeframer::new();
d.push(&wire);
assert_eq!(&d.next_message().expect("p").expect("m1")[..], b"first");
assert_eq!(&d.next_message().expect("p").expect("m2")[..], b"second");
assert_eq!(&d.next_message().expect("p").expect("m3")[..], b"third");
assert!(d.next_message().expect("p").is_none());
}
#[test]
fn round_trip_empty_message() {
let wire = frame_message(b"").expect("frame empty");
let mut d = RecordDeframer::new();
d.push(&wire);
let got = d.next_message().expect("parse").expect("empty message");
assert_eq!(got.len(), 0);
}
#[test]
fn round_trip_large_message_spans_multiple_records() {
let msg: Vec<u8> = (0..40_000u32).map(|i| (i % 251) as u8).collect();
let wire = frame_message(&msg).expect("frame");
let record_count = count_records(&wire);
assert!(record_count >= 3, "expected ≥3 records, got {record_count}");
let mut d = RecordDeframer::new();
d.set_phase(FramePhase::Established); d.push(&wire);
let got = d.next_message().expect("parse").expect("message");
assert_eq!(got.len(), msg.len());
assert_eq!(&got[..], &msg[..]);
}
#[test]
fn dribble_one_byte_at_a_time_reassembles() {
let msg = b"dribbled across many tiny reads";
let wire = frame_message(msg).expect("frame");
let mut d = RecordDeframer::new();
for b in &wire {
assert!(
d.next_message()
.expect("parse must not error mid-dribble")
.is_none(),
"no message should complete before the last byte is pushed"
);
d.push(std::slice::from_ref(b));
}
let got = d
.next_message()
.expect("parse")
.expect("message after full dribble");
assert_eq!(&got[..], msg);
}
#[test]
fn oversized_declared_record_rejected() {
let wire = [0x17u8, 0x03, 0x03, 0xFF, 0xFF, 0x00, 0x00, 0x00];
let mut d = RecordDeframer::new();
d.push(&wire);
let err = d
.next_message()
.expect_err("must reject oversized declared record");
assert!(matches!(err, CoreError::NetworkError(_)));
}
#[test]
fn oversized_inner_message_rejected_in_handshake_phase() {
let big = vec![0u8; 100 * 1024]; let wire = frame_message(&big).expect("frame");
let mut d = RecordDeframer::new(); d.push(&wire);
let err = d
.next_message()
.expect_err("must reject an oversized inner message in handshake phase");
assert!(matches!(err, CoreError::NetworkError(_)));
}
#[test]
fn established_phase_accepts_larger_message() {
let big = vec![0x5au8; 100 * 1024]; let wire = frame_message(&big).expect("frame");
let mut d = RecordDeframer::new();
d.set_phase(FramePhase::Established);
d.push(&wire);
let got = d.next_message().expect("parse").expect("message");
assert_eq!(got.len(), big.len());
}
#[test]
fn empty_record_flood_is_bounded() {
let mut d = RecordDeframer::new();
let empty_record = [CT_APPLICATION_DATA, 0x03, 0x03, 0x00, 0x00];
for _ in 0..(MAX_CONSECUTIVE_EMPTY + 2) {
d.push(&empty_record);
}
let err = d
.next_message()
.expect_err("must reject an empty-record flood");
assert!(matches!(err, CoreError::NetworkError(_)));
}
#[test]
fn encode_rejects_oversized_fragment() {
let huge = vec![0u8; 70_000]; let mut out = Vec::new();
let err = encode_record(CT_APPLICATION_DATA, VER_TLS12, &huge, &mut out)
.expect_err("must reject an oversized fragment");
assert!(matches!(err, CoreError::NetworkError(_)));
assert!(out.is_empty(), "nothing written on the error path");
}
#[test]
fn non_application_data_content_type_rejected() {
let mut wire = Vec::new();
encode_record(CT_HANDSHAKE, VER_TLS12, b"\x00\x05hello", &mut wire).expect("encode");
let mut d = RecordDeframer::new();
d.push(&wire);
let err = d
.next_message()
.expect_err("must reject 0x16 in the data phase");
assert!(matches!(err, CoreError::NetworkError(_)));
}
#[test]
fn padded_record_strips_padding() {
let mut stream = Vec::new();
stream.extend_from_slice(&3u32.to_be_bytes());
stream.extend_from_slice(b"PAD");
let wire = padded_record(&stream, 500);
assert!(wire.len() > 500);
let mut d = RecordDeframer::new();
d.push(&wire);
let got = d.next_message().expect("parse").expect("message");
assert_eq!(&got[..], b"PAD");
}
#[test]
fn inner_chunk_len_overrun_rejected() {
let mut wire = Vec::new();
let fragment = [0x00u8, 0x10, b'a', b'b', b'c'];
encode_record(CT_APPLICATION_DATA, VER_TLS12, &fragment, &mut wire).expect("encode");
let mut d = RecordDeframer::new();
d.push(&wire);
let err = d
.next_message()
.expect_err("must reject inner chunk_len overrun");
assert!(matches!(err, CoreError::NetworkError(_)));
}
#[test]
fn reassembly_buffer_is_bounded_by_cap() {
let mut d = RecordDeframer::new();
let chunk = vec![0u8; 16_000];
let mut wire = Vec::new();
for _ in 0..25 {
let mut frag = Vec::new();
frag.extend_from_slice(&(chunk.len() as u16).to_be_bytes());
frag.extend_from_slice(&chunk);
encode_record(CT_APPLICATION_DATA, VER_TLS12, &frag, &mut wire).expect("encode");
}
d.push(&wire);
let err = d
.next_message()
.expect_err("the reassembly buffer must be bounded by cap, not grow to the fed size");
assert!(matches!(err, CoreError::NetworkError(_)));
}
#[test]
fn msg_len_prefix_split_across_records() {
let mut stream = Vec::new();
stream.extend_from_slice(&2u32.to_be_bytes());
stream.extend_from_slice(b"hi");
let mut wire = Vec::new();
for part in [&stream[0..3], &stream[3..6]] {
let mut frag = Vec::new();
frag.extend_from_slice(&(part.len() as u16).to_be_bytes());
frag.extend_from_slice(part);
encode_record(CT_APPLICATION_DATA, VER_TLS12, &frag, &mut wire).expect("encode");
}
let mut d = RecordDeframer::new();
d.push(&wire);
let got = d
.next_message()
.expect("parse")
.expect("message across split prefix");
assert_eq!(&got[..], b"hi");
}
fn count_records(mut buf: &[u8]) -> usize {
let mut n = 0;
while buf.len() >= TLS_RECORD_HEADER_LEN {
let frag_len = u16::from_be_bytes([buf[3], buf[4]]) as usize;
let total = TLS_RECORD_HEADER_LEN + frag_len;
if buf.len() < total {
break;
}
n += 1;
buf = &buf[total..];
}
n
}
}