use crate::noise::Error;
pub const MAGIC: [u8; 5] = [0xff, 0xff, 0xff, b'T', b'S'];
pub const PROBE_LEN: usize = MAGIC.len() + 4;
pub const MAX_PAYLOAD: usize = 1024;
#[derive(Debug, PartialEq, Eq)]
pub enum Early<'a> {
Payload(&'a [u8]),
PushBack(&'a [u8; PROBE_LEN]),
}
pub struct EarlyReader {
buffer: [u8; PROBE_LEN + MAX_PAYLOAD],
filled: usize,
payload_len: Option<usize>,
}
impl Default for EarlyReader {
fn default() -> Self {
Self::new()
}
}
impl EarlyReader {
pub fn new() -> Self {
Self {
buffer: [0; PROBE_LEN + MAX_PAYLOAD],
filled: 0,
payload_len: None,
}
}
pub fn push(&mut self, bytes: &[u8]) -> Result<usize, Error> {
let mut consumed = 0;
loop {
let wanted = self.wanted()?;
if wanted == 0 || consumed == bytes.len() {
return Ok(consumed);
}
let take = (bytes.len() - consumed).min(wanted);
self.buffer[self.filled..self.filled + take]
.copy_from_slice(&bytes[consumed..consumed + take]);
self.filled += take;
consumed += take;
self.resolve_length()?;
}
}
fn resolve_length(&mut self) -> Result<(), Error> {
if self.payload_len.is_some() || self.filled < PROBE_LEN || !self.magic_matches() {
return Ok(());
}
let len = u32::from_be_bytes([
self.buffer[5],
self.buffer[6],
self.buffer[7],
self.buffer[8],
]) as usize;
if len > MAX_PAYLOAD {
return Err(Error::RecordTooLong);
}
self.payload_len = Some(len);
Ok(())
}
fn wanted(&self) -> Result<usize, Error> {
match self.payload_len {
None => Ok(PROBE_LEN.saturating_sub(self.filled)),
Some(len) => Ok((PROBE_LEN + len).saturating_sub(self.filled)),
}
}
pub fn is_complete(&self) -> bool {
self.wanted().is_ok_and(|w| w == 0)
}
fn magic_matches(&self) -> bool {
self.buffer[..MAGIC.len()] == MAGIC
}
pub fn finish(&self) -> Result<Early<'_>, Error> {
if !self.is_complete() {
return Err(Error::Incomplete);
}
match self.payload_len {
Some(len) => Ok(Early::Payload(&self.buffer[PROBE_LEN..PROBE_LEN + len])),
None => Ok(Early::PushBack(
self.buffer[..PROBE_LEN]
.try_into()
.expect("completeness implies the probe is full"),
)),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn captured_shape() -> ([u8; PROBE_LEN], &'static [u8]) {
let payload = br#"{"nodeKeyChallenge":"nodekey:0000000000000000000000000000000000000000000000000000000000000000"}"#;
let mut probe = [0u8; PROBE_LEN];
probe[..5].copy_from_slice(&MAGIC);
probe[5..].copy_from_slice(&(payload.len() as u32).to_be_bytes());
(probe, payload)
}
#[test]
fn reads_an_early_payload_delivered_all_at_once() {
let (probe, payload) = captured_shape();
let mut reader = EarlyReader::new();
reader.push(&probe).unwrap();
reader.push(payload).unwrap();
assert_eq!(reader.finish(), Ok(Early::Payload(payload)));
}
#[test]
fn reads_an_early_payload_split_across_record_boundaries() {
let (probe, payload) = captured_shape();
let mut reader = EarlyReader::new();
assert_eq!(reader.push(&probe[..5]).unwrap(), 5);
assert!(!reader.is_complete());
assert_eq!(reader.finish(), Err(Error::Incomplete));
assert_eq!(reader.push(&probe[5..9]).unwrap(), 4);
assert!(!reader.is_complete());
for chunk in payload.chunks(7) {
reader.push(chunk).unwrap();
}
assert_eq!(reader.finish(), Ok(Early::Payload(payload)));
}
#[test]
fn an_http2_settings_frame_is_pushed_back_intact() {
let header = [0x00, 0x00, 0x12, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00];
let mut reader = EarlyReader::new();
assert_eq!(reader.push(&header).unwrap(), PROBE_LEN);
assert_eq!(reader.finish(), Ok(Early::PushBack(&header)));
}
#[test]
fn a_probe_that_only_partly_matches_the_magic_is_pushed_back() {
let mut header = [0u8; PROBE_LEN];
header[..4].copy_from_slice(&MAGIC[..4]);
header[4] = b'X';
let mut reader = EarlyReader::new();
reader.push(&header).unwrap();
assert_eq!(reader.finish(), Ok(Early::PushBack(&header)));
}
#[test]
fn push_reports_what_it_consumed_so_the_rest_reaches_http2() {
let (probe, payload) = captured_shape();
let mut stream = [0u8; PROBE_LEN + 200];
stream[..PROBE_LEN].copy_from_slice(&probe);
stream[PROBE_LEN..PROBE_LEN + payload.len()].copy_from_slice(payload);
let total = PROBE_LEN + payload.len();
stream[total..total + 9].copy_from_slice(&[0, 0, 0x12, 4, 0, 0, 0, 0, 0]);
let mut reader = EarlyReader::new();
let consumed = reader.push(&stream[..total + 9]).unwrap();
assert_eq!(consumed, total, "the HTTP/2 frame must not be swallowed");
assert_eq!(reader.finish(), Ok(Early::Payload(payload)));
}
#[test]
fn a_length_larger_than_the_bound_is_refused() {
let mut probe = [0u8; PROBE_LEN];
probe[..5].copy_from_slice(&MAGIC);
probe[5..].copy_from_slice(&(MAX_PAYLOAD as u32 + 1).to_be_bytes());
let mut reader = EarlyReader::new();
assert_eq!(reader.push(&probe), Err(Error::RecordTooLong));
}
}