watermelon_proto/proto/decoder/

stream.rs

use bytes::{BufMut, BytesMut};

use crate::{
    proto::{ServerOp, error::DecoderError},
    util::CrlfFinder,
};

use super::DecoderStatus;

const INITIAL_READ_BUF_CAPACITY: usize = 64 * 1024;

#[derive(Debug)]
pub struct StreamDecoder {
    read_buf: BytesMut,
    status: DecoderStatus,
    crlf: CrlfFinder,
}

impl StreamDecoder {
    #[must_use]
    pub fn new() -> Self {
        Self {
            read_buf: BytesMut::with_capacity(INITIAL_READ_BUF_CAPACITY),
            status: DecoderStatus::ControlLine { last_bytes_read: 0 },
            crlf: CrlfFinder::new(),
        }
    }

    #[must_use]
    pub fn read_buf(&mut self) -> &mut (impl BufMut + use<>) {
        &mut self.read_buf
    }

    /// Decodes the next frame of bytes into a [`ServerOp`].
    ///
    /// A `None` variant is returned in case no progress is made,
    ///
    /// # Errors
    ///
    /// It returns an error if a decoding error occurs.
    pub fn decode(&mut self) -> Result<Option<ServerOp>, DecoderError> {
        super::decode(&self.crlf, &mut self.status, &mut self.read_buf)
    }
}

impl Default for StreamDecoder {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use bytes::{BufMut as _, Bytes};
    use claims::{assert_matches, assert_ok_eq};

    use bytestring::ByteString;

    use crate::{
        StatusCode, Subject,
        error::ServerError,
        headers::{HeaderMap, HeaderName, HeaderValue},
        message::{MessageBase, ServerMessage},
        proto::{error::DecoderError, server::ServerOp},
        util::ParseUintError,
    };

    use super::StreamDecoder;

    #[test]
    fn decode_ping() {
        let mut decoder = StreamDecoder::new();
        decoder.read_buf().put_slice(b"PING\r\n");
        assert_ok_eq!(decoder.decode(), Some(ServerOp::Ping));
        assert_ok_eq!(decoder.decode(), None);
    }

    #[test]
    fn decode_pong() {
        let mut decoder = StreamDecoder::new();
        decoder.read_buf().put_slice(b"PONG\r\n");
        assert_ok_eq!(decoder.decode(), Some(ServerOp::Pong));
        assert_ok_eq!(decoder.decode(), None);
    }

    #[test]
    fn decode_ok() {
        let mut decoder = StreamDecoder::new();
        decoder.read_buf().put_slice(b"+OK\r\n");
        assert_ok_eq!(decoder.decode(), Some(ServerOp::Success));
        assert_ok_eq!(decoder.decode(), None);
    }

    #[test]
    fn decode_error() {
        let mut decoder = StreamDecoder::new();
        decoder
            .read_buf()
            .put_slice(b"-ERR 'Authorization Violation'\r\n");
        assert_ok_eq!(
            decoder.decode(),
            Some(ServerOp::Error {
                error: ServerError::AuthorizationViolation
            })
        );
        assert_ok_eq!(decoder.decode(), None);
    }

    #[test]
    fn decode_error_bare_quote() {
        let mut decoder = StreamDecoder::new();
        // A single-quote-only message used to underflow the quote
        // stripping logic (`len() - 1` on an empty buffer).
        decoder.read_buf().put_slice(b"-ERR '\r\n");
        assert_matches!(decoder.decode(), Err(DecoderError::InvalidErrorMessage));
    }

    #[test]
    fn decode_error_empty_message() {
        let mut decoder = StreamDecoder::new();
        decoder.read_buf().put_slice(b"-ERR ''\r\n");
        assert_ok_eq!(
            decoder.decode(),
            Some(ServerOp::Error {
                error: ServerError::Other {
                    raw_message: ByteString::new()
                }
            })
        );
    }

    #[test]
    fn decode_msg() {
        let mut decoder = StreamDecoder::new();
        decoder
            .read_buf()
            .put_slice(b"MSG hello.world 1 12\r\nHello World!\r\n");
        assert_ok_eq!(
            decoder.decode(),
            Some(ServerOp::Message {
                message: ServerMessage {
                    status_code: None,
                    status_description: None,
                    subscription_id: 1.into(),
                    base: MessageBase {
                        subject: Subject::from_static("hello.world"),
                        reply_subject: None,
                        headers: HeaderMap::new(),
                        payload: Bytes::from_static(b"Hello World!")
                    }
                }
            })
        );
        assert_ok_eq!(decoder.decode(), None);
    }

    #[test]
    fn decode_hmsg_status_description() {
        const HEADERS: &str = "NATS/1.0 409 Batch Completed\r\nNats-Pending-Messages: 3\r\nNats-Pending-Bytes: 0\r\n\r\n";

        let mut decoder = StreamDecoder::new();
        decoder.read_buf().put_slice(
            format!(
                "HMSG _INBOX.abcd 9 {} {}\r\n{HEADERS}\r\n",
                HEADERS.len(),
                HEADERS.len()
            )
            .as_bytes(),
        );
        assert_ok_eq!(
            decoder.decode(),
            Some(ServerOp::Message {
                message: ServerMessage {
                    status_code: Some(StatusCode::CONFLICT),
                    status_description: Some(ByteString::from_static("Batch Completed")),
                    subscription_id: 9.into(),
                    base: MessageBase {
                        subject: Subject::from_static("_INBOX.abcd"),
                        reply_subject: None,
                        headers: HeaderMap::from_iter([
                            (
                                HeaderName::from_static("Nats-Pending-Messages"),
                                HeaderValue::from_static("3")
                            ),
                            (
                                HeaderName::from_static("Nats-Pending-Bytes"),
                                HeaderValue::from_static("0")
                            ),
                        ]),
                        payload: Bytes::new(),
                    }
                }
            })
        );
        assert_ok_eq!(decoder.decode(), None);
    }

    #[test]
    fn decode_hmsg_status_without_description() {
        const HEADERS: &str = "NATS/1.0 100\r\n\r\n";

        let mut decoder = StreamDecoder::new();
        decoder.read_buf().put_slice(
            format!(
                "HMSG _INBOX.abcd 9 {} {}\r\n{HEADERS}\r\n",
                HEADERS.len(),
                HEADERS.len()
            )
            .as_bytes(),
        );
        assert_ok_eq!(
            decoder.decode(),
            Some(ServerOp::Message {
                message: ServerMessage {
                    status_code: Some(StatusCode::IDLE_HEARTBEAT),
                    status_description: None,
                    subscription_id: 9.into(),
                    base: MessageBase {
                        subject: Subject::from_static("_INBOX.abcd"),
                        reply_subject: None,
                        headers: HeaderMap::new(),
                        payload: Bytes::new(),
                    }
                }
            })
        );
        assert_ok_eq!(decoder.decode(), None);
    }

    #[test]
    fn decode_hmsg_status_invalid_description_separator() {
        const HEADERS: &str = "NATS/1.0 409XBatch Completed\r\n\r\n";

        let mut decoder = StreamDecoder::new();
        decoder.read_buf().put_slice(
            format!(
                "HMSG _INBOX.abcd 9 {} {}\r\n{HEADERS}\r\n",
                HEADERS.len(),
                HEADERS.len()
            )
            .as_bytes(),
        );
        assert_matches!(decoder.decode(), Err(DecoderError::InvalidHead));
    }

    #[test]
    fn decode_hmsg_non_utf8_header_value() {
        const HEADERS: &[u8] = b"NATS/1.0\r\nX-Binary: \x00\x80\xffabc\r\n\r\n";

        let mut decoder = StreamDecoder::new();
        decoder.read_buf().put_slice(
            format!("HMSG _INBOX.abcd 9 {} {}\r\n", HEADERS.len(), HEADERS.len()).as_bytes(),
        );
        decoder.read_buf().put_slice(HEADERS);
        decoder.read_buf().put_slice(b"\r\n");
        assert_ok_eq!(
            decoder.decode(),
            Some(ServerOp::Message {
                message: ServerMessage {
                    status_code: None,
                    status_description: None,
                    subscription_id: 9.into(),
                    base: MessageBase {
                        subject: Subject::from_static("_INBOX.abcd"),
                        reply_subject: None,
                        headers: HeaderMap::from_iter([(
                            HeaderName::from_static("X-Binary"),
                            HeaderValue::from_bytes(b"\x00\x80\xffabc").unwrap(),
                        )]),
                        payload: Bytes::new(),
                    }
                }
            })
        );
        assert_ok_eq!(decoder.decode(), None);
    }

    #[test]
    fn decode_hmsg_empty_header_value() {
        const HEADERS: &str = "NATS/1.0\r\nA:\r\nB: \r\n\r\n";

        let mut decoder = StreamDecoder::new();
        decoder.read_buf().put_slice(
            format!(
                "HMSG _INBOX.abcd 9 {} {}\r\n{HEADERS}\r\n",
                HEADERS.len(),
                HEADERS.len()
            )
            .as_bytes(),
        );
        assert_ok_eq!(
            decoder.decode(),
            Some(ServerOp::Message {
                message: ServerMessage {
                    status_code: None,
                    status_description: None,
                    subscription_id: 9.into(),
                    base: MessageBase {
                        subject: Subject::from_static("_INBOX.abcd"),
                        reply_subject: None,
                        headers: HeaderMap::from_iter([
                            (
                                HeaderName::from_static("A"),
                                HeaderValue::from_bytes(b"").unwrap(),
                            ),
                            (
                                HeaderName::from_static("B"),
                                HeaderValue::from_bytes(b"").unwrap(),
                            ),
                        ]),
                        payload: Bytes::new(),
                    }
                }
            })
        );
        assert_ok_eq!(decoder.decode(), None);
    }

    #[test]
    fn msg_huge_payload_len() {
        let mut decoder = StreamDecoder::new();
        decoder
            .read_buf()
            .put_slice(b"MSG hello.world 1 18446744073709551615\r\nHello");
        // The decoder must reject the impossibly long declared payload
        // instead of waiting for it (or, historically, overflowing while
        // checking the buffered length and panicking in `split_to`).
        assert_matches!(
            decoder.decode(),
            Err(DecoderError::InvalidPayloadLength(ParseUintError::Overflow))
        );
    }

    #[test]
    fn msg_payload_len_over_limit() {
        let mut decoder = StreamDecoder::new();
        decoder
            .read_buf()
            .put_slice(b"MSG hello.world 1 134217729\r\n");
        assert_matches!(
            decoder.decode(),
            Err(DecoderError::InvalidPayloadLength(ParseUintError::Overflow))
        );
    }

    #[test]
    fn hmsg_total_len_over_limit() {
        let mut decoder = StreamDecoder::new();
        decoder
            .read_buf()
            .put_slice(b"HMSG hello.world 1 100 134217729\r\n");
        assert_matches!(
            decoder.decode(),
            Err(DecoderError::InvalidPayloadLength(ParseUintError::Overflow))
        );
    }

    #[test]
    fn head_too_long() {
        let mut decoder = StreamDecoder::new();
        decoder.read_buf().put_bytes(0, 200_000);
        assert_matches!(
            decoder.decode(),
            Err(DecoderError::HeadTooLong { len: 200_000 })
        );
    }
}