WEB

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
Co-Authored-By: John Preston <17900494+john-preston@users.noreply.github.com>

src/web/frame.rs

@@ -0,0 +1,272 @@
+use bytes::{BufMut, Bytes, BytesMut};
+
+use crate::config::WebLimitsConfig;
+
+/// Fixed WEB frame header size.
+pub(crate) const HEADER_BYTES: usize = 8;
+/// Initial bidirectional stream credit.
+pub(crate) const INITIAL_STREAM_WINDOW: u32 = 4 * 1024 * 1024;
+/// Maximum data chunk emitted by the server.
+pub(crate) const DATA_CHUNK_BYTES: usize = 64 * 1024;
+
+/// WEB frame type codes shared with Telegram Desktop.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[repr(u8)]
+pub(crate) enum FrameType {
+    /// Opens a logical MTProxy stream.
+    Open = 0x01,
+    /// Carries logical-stream payload bytes.
+    Data = 0x02,
+    /// Closes a logical stream.
+    Close = 0x03,
+    /// Returns consumed flow-control credit.
+    Window = 0x04,
+    /// Requests an application-level liveness response.
+    Ping = 0x05,
+    /// Answers application-level liveness traffic.
+    Pong = 0x06,
+    /// Starts one WEB carrier session.
+    Hello = 0x10,
+    /// Confirms WEB carrier session creation.
+    Welcome = 0x11,
+    /// Terminates a WEB carrier session.
+    Bye = 0x1f,
+}
+
+impl FrameType {
+    fn parse(value: u8) -> Option<Self> {
+        match value {
+            0x01 => Some(Self::Open),
+            0x02 => Some(Self::Data),
+            0x03 => Some(Self::Close),
+            0x04 => Some(Self::Window),
+            0x05 => Some(Self::Ping),
+            0x06 => Some(Self::Pong),
+            0x10 => Some(Self::Hello),
+            0x11 => Some(Self::Welcome),
+            0x1f => Some(Self::Bye),
+            _ => None,
+        }
+    }
+}
+
+/// One parsed frame borrowing its payload from the HTTP request body.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) struct Frame<'a> {
+    /// Parsed frame type.
+    pub(crate) frame_type: FrameType,
+    /// Logical 24-bit stream identifier.
+    pub(crate) stream_id: u32,
+    /// Borrowed frame payload.
+    pub(crate) payload: &'a [u8],
+}
+
+/// Protocol parse or shape failure.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum FrameError {
+    /// A carrier body contained no frame.
+    EmptyBatch,
+    /// A carrier body exceeded the configured frame count.
+    TooManyFrames,
+    /// A frame header or payload was truncated.
+    Incomplete,
+    /// A frame payload exceeded its configured ceiling.
+    PayloadLimit,
+    /// The frame type code is not defined.
+    UnknownType,
+    /// A known frame violated direction-specific grammar.
+    InvalidShape,
+}
+
+/// Parses and validates all frame boundaries without copying payloads.
+pub(crate) fn parse_all<'a>(
+    input: &'a [u8],
+    limits: &WebLimitsConfig,
+) -> std::result::Result<Vec<Frame<'a>>, FrameError> {
+    if input.is_empty() {
+        return Err(FrameError::EmptyBatch);
+    }
+    let mut remaining = input;
+    let mut frames = Vec::with_capacity(remaining.len().div_ceil(HEADER_BYTES).min(16));
+    while !remaining.is_empty() {
+        if frames.len() >= limits.max_frames_per_body {
+            return Err(FrameError::TooManyFrames);
+        }
+        if remaining.len() < HEADER_BYTES {
+            return Err(FrameError::Incomplete);
+        }
+        let frame_type = FrameType::parse(remaining[0]).ok_or(FrameError::UnknownType)?;
+        let stream_id = u32::from(remaining[1]) << 16
+            | u32::from(remaining[2]) << 8
+            | u32::from(remaining[3]);
+        let payload_len = u32::from_be_bytes([
+            remaining[4],
+            remaining[5],
+            remaining[6],
+            remaining[7],
+        ]) as usize;
+        if payload_len > limits.max_frame_payload_bytes {
+            return Err(FrameError::PayloadLimit);
+        }
+        let frame_len = HEADER_BYTES
+            .checked_add(payload_len)
+            .ok_or(FrameError::PayloadLimit)?;
+        if frame_len > remaining.len() {
+            return Err(FrameError::Incomplete);
+        }
+        frames.push(Frame {
+            frame_type,
+            stream_id,
+            payload: &remaining[HEADER_BYTES..frame_len],
+        });
+        remaining = &remaining[frame_len..];
+    }
+    Ok(frames)
+}
+
+/// Enforces the client-to-server frame grammar.
+pub(crate) fn validate_client_shape(frame: Frame<'_>) -> std::result::Result<(), FrameError> {
+    if frame.stream_id == 0 {
+        return if frame.frame_type == FrameType::Pong && frame.payload.len() <= 64 {
+            Ok(())
+        } else {
+            Err(FrameError::InvalidShape)
+        };
+    }
+    match frame.frame_type {
+        FrameType::Open | FrameType::Close if frame.payload.is_empty() => Ok(()),
+        FrameType::Data if !frame.payload.is_empty() => Ok(()),
+        FrameType::Window => window_amount(frame.payload).map(|_| ()),
+        _ => Err(FrameError::InvalidShape),
+    }
+}
+
+/// Validates the exact first-session HELLO body.
+pub(crate) fn validate_hello(input: &[u8], limits: &WebLimitsConfig) -> bool {
+    let Ok(frames) = parse_all(input, limits) else {
+        return false;
+    };
+    frames.len() == 1
+        && frames[0].frame_type == FrameType::Hello
+        && frames[0].stream_id == 0
+        && frames[0].payload == [1]
+}
+
+/// Encodes one complete WEB frame.
+pub(crate) fn encode(frame_type: FrameType, stream_id: u32, payload: &[u8]) -> Bytes {
+    let mut output = BytesMut::with_capacity(HEADER_BYTES + payload.len());
+    output.put_u8(frame_type as u8);
+    output.put_u8((stream_id >> 16) as u8);
+    output.put_u8((stream_id >> 8) as u8);
+    output.put_u8(stream_id as u8);
+    output.put_u32(payload.len() as u32);
+    output.extend_from_slice(payload);
+    output.freeze()
+}
+
+/// Decodes a non-zero WINDOW delta.
+pub(crate) fn window_amount(payload: &[u8]) -> std::result::Result<u32, FrameError> {
+    let bytes: [u8; 4] = payload.try_into().map_err(|_| FrameError::InvalidShape)?;
+    let amount = u32::from_be_bytes(bytes);
+    (amount != 0)
+        .then_some(amount)
+        .ok_or(FrameError::InvalidShape)
+}
+
+/// Encodes a WINDOW delta payload.
+pub(crate) fn window_payload(amount: u32) -> [u8; 4] {
+    amount.to_be_bytes()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn hello_and_welcome_match_reference_bytes() {
+        let limits = WebLimitsConfig::default();
+        let hello = encode(FrameType::Hello, 0, &[1]);
+        assert_eq!(hello.as_ref(), &hex::decode("100000000000000101").unwrap());
+        assert!(validate_hello(&hello, &limits));
+        assert_eq!(
+            encode(FrameType::Welcome, 0, &[]).as_ref(),
+            &hex::decode("1100000000000000").unwrap()
+        );
+    }
+
+    #[test]
+    fn parser_rejects_excessive_payload_before_slicing() {
+        let limits = WebLimitsConfig {
+            max_frame_payload_bytes: 4,
+            ..WebLimitsConfig::default()
+        };
+        let frame = encode(FrameType::Data, 1, &[0; 5]);
+        assert_eq!(parse_all(&frame, &limits), Err(FrameError::PayloadLimit));
+    }
+
+    #[test]
+    fn client_shape_rejects_control_types_on_stream_zero() {
+        let frame = Frame {
+            frame_type: FrameType::Ping,
+            stream_id: 0,
+            payload: &[],
+        };
+        assert_eq!(validate_client_shape(frame), Err(FrameError::InvalidShape));
+    }
+
+    #[test]
+    fn stream_frames_match_client_reference_vectors() {
+        assert_eq!(
+            encode(FrameType::Open, 17, &[]).as_ref(),
+            &hex::decode("0100001100000000").unwrap()
+        );
+        assert_eq!(
+            encode(FrameType::Data, 17, b"round trip").as_ref(),
+            &hex::decode("020000110000000a726f756e642074726970").unwrap()
+        );
+        assert_eq!(
+            encode(FrameType::Window, 17, &10u32.to_be_bytes()).as_ref(),
+            &hex::decode("04000011000000040000000a").unwrap()
+        );
+        assert_eq!(
+            encode(FrameType::Open, 0x00ff_ffff, &[]).as_ref(),
+            &hex::decode("01ffffff00000000").unwrap()
+        );
+    }
+
+    #[test]
+    fn parser_rejects_empty_truncated_and_excessive_batches() {
+        let mut limits = WebLimitsConfig::default();
+        assert_eq!(parse_all(&[], &limits), Err(FrameError::EmptyBatch));
+        assert_eq!(
+            parse_all(&hex::decode("0200000100000001").unwrap(), &limits),
+            Err(FrameError::Incomplete)
+        );
+        limits.max_frames_per_body = 1;
+        let mut body = encode(FrameType::Pong, 0, &[]).to_vec();
+        body.extend_from_slice(&encode(FrameType::Pong, 0, &[]));
+        assert_eq!(parse_all(&body, &limits), Err(FrameError::TooManyFrames));
+    }
+
+    #[test]
+    fn client_shape_rejects_empty_data_and_zero_window() {
+        let empty_data = Frame {
+            frame_type: FrameType::Data,
+            stream_id: 1,
+            payload: &[],
+        };
+        let zero_window = Frame {
+            frame_type: FrameType::Window,
+            stream_id: 1,
+            payload: &[0; 4],
+        };
+        assert_eq!(
+            validate_client_shape(empty_data),
+            Err(FrameError::InvalidShape)
+        );
+        assert_eq!(
+            validate_client_shape(zero_window),
+            Err(FrameError::InvalidShape)
+        );
+    }
+}