use crate::enums::ContentType;
use crate::enums::ProtocolVersion;
use crate::msgs::message::{BorrowedPlainMessage, PlainMessage};
use crate::Error;
pub const MAX_FRAGMENT_LEN: usize = 16384;
pub const PACKET_OVERHEAD: usize = 1 + 2 + 2;
pub const MAX_FRAGMENT_SIZE: usize = MAX_FRAGMENT_LEN + PACKET_OVERHEAD;
pub struct MessageFragmenter {
max_frag: usize,
}
impl Default for MessageFragmenter {
fn default() -> Self {
Self {
max_frag: MAX_FRAGMENT_LEN,
}
}
}
impl MessageFragmenter {
pub fn fragment_message<'a>(
&self,
msg: &'a PlainMessage,
) -> impl Iterator<Item = BorrowedPlainMessage<'a>> + 'a {
self.fragment_slice(msg.typ, msg.version, &msg.payload.0)
}
pub fn fragment_slice<'a>(
&self,
typ: ContentType,
version: ProtocolVersion,
payload: &'a [u8],
) -> impl Iterator<Item = BorrowedPlainMessage<'a>> + 'a {
payload
.chunks(self.max_frag)
.map(move |c| BorrowedPlainMessage {
typ,
version,
payload: c,
})
}
pub fn set_max_fragment_size(&mut self, max_fragment_size: Option<usize>) -> Result<(), Error> {
self.max_frag = match max_fragment_size {
Some(sz @ 32..=MAX_FRAGMENT_SIZE) => sz - PACKET_OVERHEAD,
None => MAX_FRAGMENT_LEN,
_ => return Err(Error::BadMaxFragmentSize),
};
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::{MessageFragmenter, PACKET_OVERHEAD};
use crate::enums::ContentType;
use crate::enums::ProtocolVersion;
use crate::msgs::base::Payload;
use crate::msgs::message::{BorrowedPlainMessage, PlainMessage};
fn msg_eq(
m: &BorrowedPlainMessage,
total_len: usize,
typ: &ContentType,
version: &ProtocolVersion,
bytes: &[u8],
) {
assert_eq!(&m.typ, typ);
assert_eq!(&m.version, version);
assert_eq!(m.payload, bytes);
let buf = m.to_unencrypted_opaque().encode();
assert_eq!(total_len, buf.len());
}
#[test]
fn smoke() {
let typ = ContentType::Handshake;
let version = ProtocolVersion::TLSv1_2;
let data: Vec<u8> = (1..70u8).collect();
let m = PlainMessage {
typ,
version,
payload: Payload::new(data),
};
let mut frag = MessageFragmenter::default();
frag.set_max_fragment_size(Some(32))
.unwrap();
let q = frag
.fragment_message(&m)
.collect::<Vec<_>>();
assert_eq!(q.len(), 3);
msg_eq(
&q[0],
32,
&typ,
&version,
&[
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23,
24, 25, 26, 27,
],
);
msg_eq(
&q[1],
32,
&typ,
&version,
&[
28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48,
49, 50, 51, 52, 53, 54,
],
);
msg_eq(
&q[2],
20,
&typ,
&version,
&[55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69],
);
}
#[test]
fn non_fragment() {
let m = PlainMessage {
typ: ContentType::Handshake,
version: ProtocolVersion::TLSv1_2,
payload: Payload::new(b"\x01\x02\x03\x04\x05\x06\x07\x08".to_vec()),
};
let mut frag = MessageFragmenter::default();
frag.set_max_fragment_size(Some(32))
.unwrap();
let q = frag
.fragment_message(&m)
.collect::<Vec<_>>();
assert_eq!(q.len(), 1);
msg_eq(
&q[0],
PACKET_OVERHEAD + 8,
&ContentType::Handshake,
&ProtocolVersion::TLSv1_2,
b"\x01\x02\x03\x04\x05\x06\x07\x08",
);
}
}