ember_protocol/

parse.rs

//! Zero-copy RESP3 parser.
//!
//! Operates on buffered byte slices. The caller is responsible for reading
//! data from the network into a buffer — this parser is purely synchronous.
//!
//! The parser uses a `Cursor<&[u8]>` to track its position through the
//! input buffer without consuming it, allowing the caller to retry once
//! more data arrives.
//!
//! # Single-pass design
//!
//! Earlier versions used a two-pass approach: `check()` to validate a
//! complete frame exists, then `parse()` to build Frame values. This
//! scanned every byte twice. The current implementation does a single
//! pass that builds Frame values directly, returning `Incomplete` if
//! the buffer doesn't contain enough data yet.

use std::io::Cursor;

use bytes::Bytes;

use crate::error::ProtocolError;
use crate::types::Frame;

/// Maximum nesting depth for arrays and maps. Prevents stack overflow
/// from malicious or malformed deeply-nested frames.
const MAX_NESTING_DEPTH: usize = 64;

/// Maximum number of elements in an array or map. Prevents memory
/// amplification attacks where tiny elements (3 bytes each) create
/// disproportionately large Vec allocations.
const MAX_ARRAY_ELEMENTS: usize = 1_048_576;

/// Maximum length of a bulk string in bytes (512 MB, matching Redis).
const MAX_BULK_LEN: i64 = 512 * 1024 * 1024;

/// Cap for Vec::with_capacity in array/map parsing. A declared count of
/// 1M elements with capacity pre-allocation costs ~72 MB upfront even
/// before any child data is parsed. This cap limits the initial allocation
/// while still letting the Vec grow organically as elements are parsed.
const PREALLOC_CAP: usize = 1024;

/// Checks whether `buf` contains a complete RESP3 frame and parses it.
///
/// Returns `Ok(Some(frame))` if a complete frame was parsed,
/// `Ok(None)` if the buffer doesn't contain enough data yet,
/// or `Err(...)` if the data is malformed.
#[inline]
pub fn parse_frame(buf: &[u8]) -> Result<Option<(Frame, usize)>, ProtocolError> {
    if buf.is_empty() {
        return Ok(None);
    }

    let mut cursor = Cursor::new(buf);

    match try_parse(&mut cursor, 0) {
        Ok(frame) => {
            let consumed = cursor.position() as usize;
            Ok(Some((frame, consumed)))
        }
        Err(ProtocolError::Incomplete) => Ok(None),
        Err(e) => Err(e),
    }
}

// ---------------------------------------------------------------------------
// single-pass parser: validates and builds Frame values in one traversal
// ---------------------------------------------------------------------------

/// Parses a complete RESP3 frame from the cursor position, returning
/// `Incomplete` if the buffer doesn't contain enough data.
fn try_parse(cursor: &mut Cursor<&[u8]>, depth: usize) -> Result<Frame, ProtocolError> {
    let prefix = read_byte(cursor)?;

    match prefix {
        b'+' => {
            let line = read_line(cursor)?;
            let s = std::str::from_utf8(line).map_err(|_| {
                ProtocolError::InvalidCommandFrame("invalid utf-8 in simple string".into())
            })?;
            Ok(Frame::Simple(s.to_owned()))
        }
        b'-' => {
            let line = read_line(cursor)?;
            let s = std::str::from_utf8(line).map_err(|_| {
                ProtocolError::InvalidCommandFrame("invalid utf-8 in error string".into())
            })?;
            Ok(Frame::Error(s.to_owned()))
        }
        b':' => {
            let val = read_integer_line(cursor)?;
            Ok(Frame::Integer(val))
        }
        b'$' => {
            let len = read_integer_line(cursor)?;
            if len < 0 {
                return Err(ProtocolError::InvalidFrameLength(len));
            }
            if len > MAX_BULK_LEN {
                return Err(ProtocolError::BulkStringTooLarge(len as usize));
            }
            let len = len as usize;

            // need `len` bytes of data + \r\n
            let remaining = remaining(cursor);
            if remaining < len + 2 {
                return Err(ProtocolError::Incomplete);
            }

            let pos = cursor.position() as usize;
            let buf = cursor.get_ref();

            // verify trailing \r\n
            if buf[pos + len] != b'\r' || buf[pos + len + 1] != b'\n' {
                return Err(ProtocolError::InvalidFrameLength(len as i64));
            }

            let data = &buf[pos..pos + len];
            cursor.set_position((pos + len + 2) as u64);
            Ok(Frame::Bulk(Bytes::copy_from_slice(data)))
        }
        b'*' => {
            let next_depth = depth + 1;
            if next_depth > MAX_NESTING_DEPTH {
                return Err(ProtocolError::NestingTooDeep(MAX_NESTING_DEPTH));
            }

            let count = read_integer_line(cursor)?;
            if count < 0 {
                return Err(ProtocolError::InvalidFrameLength(count));
            }
            if count as usize > MAX_ARRAY_ELEMENTS {
                return Err(ProtocolError::TooManyElements(count as usize));
            }

            let count = count as usize;
            let mut frames = Vec::with_capacity(count.min(PREALLOC_CAP));
            for _ in 0..count {
                frames.push(try_parse(cursor, next_depth)?);
            }
            Ok(Frame::Array(frames))
        }
        b'_' => {
            // consume the trailing \r\n
            let _ = read_line(cursor)?;
            Ok(Frame::Null)
        }
        b'%' => {
            let next_depth = depth + 1;
            if next_depth > MAX_NESTING_DEPTH {
                return Err(ProtocolError::NestingTooDeep(MAX_NESTING_DEPTH));
            }

            let count = read_integer_line(cursor)?;
            if count < 0 {
                return Err(ProtocolError::InvalidFrameLength(count));
            }
            if count as usize > MAX_ARRAY_ELEMENTS {
                return Err(ProtocolError::TooManyElements(count as usize));
            }

            let count = count as usize;
            let mut pairs = Vec::with_capacity(count.min(PREALLOC_CAP));
            for _ in 0..count {
                let key = try_parse(cursor, next_depth)?;
                let val = try_parse(cursor, next_depth)?;
                pairs.push((key, val));
            }
            Ok(Frame::Map(pairs))
        }
        other => Err(ProtocolError::InvalidPrefix(other)),
    }
}

// ---------------------------------------------------------------------------
// low-level cursor helpers
// ---------------------------------------------------------------------------

fn read_byte(cursor: &mut Cursor<&[u8]>) -> Result<u8, ProtocolError> {
    let pos = cursor.position() as usize;
    if pos >= cursor.get_ref().len() {
        return Err(ProtocolError::Incomplete);
    }
    cursor.set_position((pos + 1) as u64);
    Ok(cursor.get_ref()[pos])
}

/// Returns the slice of bytes up to (but not including) the next `\r\n`,
/// and advances the cursor past the `\r\n`.
fn read_line<'a>(cursor: &mut Cursor<&'a [u8]>) -> Result<&'a [u8], ProtocolError> {
    let start = cursor.position() as usize;
    let end = find_crlf(cursor)?;
    Ok(&cursor.get_ref()[start..end])
}

/// Reads a line and parses it as an i64.
fn read_integer_line(cursor: &mut Cursor<&[u8]>) -> Result<i64, ProtocolError> {
    let line = read_line(cursor)?;
    parse_i64(line)
}

/// Finds the next `\r\n` in the buffer starting from the cursor position.
/// Returns the index of `\r` and advances the cursor past the `\n`.
fn find_crlf(cursor: &mut Cursor<&[u8]>) -> Result<usize, ProtocolError> {
    let buf = cursor.get_ref();
    let start = cursor.position() as usize;

    if start >= buf.len() {
        return Err(ProtocolError::Incomplete);
    }

    // SIMD-accelerated scan for \r, then verify \n follows.
    // memchr processes 16-32 bytes per cycle vs 1 byte in a naive loop.
    let mut pos = start;
    while let Some(offset) = memchr::memchr(b'\r', &buf[pos..]) {
        let cr = pos + offset;
        if cr + 1 < buf.len() && buf[cr + 1] == b'\n' {
            cursor.set_position((cr + 2) as u64);
            return Ok(cr);
        }
        // bare \r without \n — keep scanning past it
        pos = cr + 1;
    }

    Err(ProtocolError::Incomplete)
}

fn remaining(cursor: &Cursor<&[u8]>) -> usize {
    let len = cursor.get_ref().len();
    let pos = cursor.position() as usize;
    len.saturating_sub(pos)
}

fn parse_i64(buf: &[u8]) -> Result<i64, ProtocolError> {
    let s = std::str::from_utf8(buf).map_err(|_| ProtocolError::InvalidInteger)?;
    s.parse::<i64>().map_err(|_| ProtocolError::InvalidInteger)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn must_parse(input: &[u8]) -> Frame {
        let (frame, consumed) = parse_frame(input)
            .expect("parse should not error")
            .expect("parse should return a frame");
        assert_eq!(consumed, input.len(), "should consume entire input");
        frame
    }

    #[test]
    fn simple_string() {
        assert_eq!(must_parse(b"+OK\r\n"), Frame::Simple("OK".into()));
        assert_eq!(
            must_parse(b"+hello world\r\n"),
            Frame::Simple("hello world".into())
        );
    }

    #[test]
    fn simple_error() {
        assert_eq!(
            must_parse(b"-ERR unknown command\r\n"),
            Frame::Error("ERR unknown command".into())
        );
    }

    #[test]
    fn integer() {
        assert_eq!(must_parse(b":42\r\n"), Frame::Integer(42));
        assert_eq!(must_parse(b":0\r\n"), Frame::Integer(0));
        assert_eq!(must_parse(b":-1\r\n"), Frame::Integer(-1));
        assert_eq!(
            must_parse(b":9223372036854775807\r\n"),
            Frame::Integer(i64::MAX)
        );
        assert_eq!(
            must_parse(b":-9223372036854775808\r\n"),
            Frame::Integer(i64::MIN)
        );
    }

    #[test]
    fn bulk_string() {
        assert_eq!(
            must_parse(b"$5\r\nhello\r\n"),
            Frame::Bulk(Bytes::from_static(b"hello"))
        );
    }

    #[test]
    fn empty_bulk_string() {
        assert_eq!(
            must_parse(b"$0\r\n\r\n"),
            Frame::Bulk(Bytes::from_static(b""))
        );
    }

    #[test]
    fn bulk_string_with_binary() {
        let input = b"$4\r\n\x00\x01\x02\x03\r\n";
        assert_eq!(
            must_parse(input),
            Frame::Bulk(Bytes::copy_from_slice(&[0, 1, 2, 3]))
        );
    }

    #[test]
    fn null() {
        assert_eq!(must_parse(b"_\r\n"), Frame::Null);
    }

    #[test]
    fn array() {
        let input = b"*2\r\n+hello\r\n+world\r\n";
        assert_eq!(
            must_parse(input),
            Frame::Array(vec![
                Frame::Simple("hello".into()),
                Frame::Simple("world".into()),
            ])
        );
    }

    #[test]
    fn empty_array() {
        assert_eq!(must_parse(b"*0\r\n"), Frame::Array(vec![]));
    }

    #[test]
    fn nested_array() {
        let input = b"*2\r\n*2\r\n:1\r\n:2\r\n*2\r\n:3\r\n:4\r\n";
        assert_eq!(
            must_parse(input),
            Frame::Array(vec![
                Frame::Array(vec![Frame::Integer(1), Frame::Integer(2)]),
                Frame::Array(vec![Frame::Integer(3), Frame::Integer(4)]),
            ])
        );
    }

    #[test]
    fn array_with_null() {
        let input = b"*3\r\n+OK\r\n_\r\n:1\r\n";
        assert_eq!(
            must_parse(input),
            Frame::Array(vec![
                Frame::Simple("OK".into()),
                Frame::Null,
                Frame::Integer(1),
            ])
        );
    }

    #[test]
    fn map() {
        let input = b"%2\r\n+key1\r\n:1\r\n+key2\r\n:2\r\n";
        assert_eq!(
            must_parse(input),
            Frame::Map(vec![
                (Frame::Simple("key1".into()), Frame::Integer(1)),
                (Frame::Simple("key2".into()), Frame::Integer(2)),
            ])
        );
    }

    #[test]
    fn incomplete_returns_none() {
        assert_eq!(parse_frame(b"").unwrap(), None);
        assert_eq!(parse_frame(b"+OK").unwrap(), None);
        assert_eq!(parse_frame(b"+OK\r").unwrap(), None);
        assert_eq!(parse_frame(b"$5\r\nhel").unwrap(), None);
        assert_eq!(parse_frame(b"*2\r\n+OK\r\n").unwrap(), None);
    }

    #[test]
    fn invalid_prefix() {
        let err = parse_frame(b"~invalid\r\n").unwrap_err();
        assert_eq!(err, ProtocolError::InvalidPrefix(b'~'));
    }

    #[test]
    fn invalid_integer() {
        let err = parse_frame(b":abc\r\n").unwrap_err();
        assert_eq!(err, ProtocolError::InvalidInteger);
    }

    #[test]
    fn negative_bulk_length() {
        let err = parse_frame(b"$-1\r\n").unwrap_err();
        assert!(matches!(err, ProtocolError::InvalidFrameLength(-1)));
    }

    #[test]
    fn parse_consumes_exact_bytes() {
        // buffer contains a full frame plus trailing garbage
        let buf = b"+OK\r\ntrailing";
        let (frame, consumed) = parse_frame(buf).unwrap().unwrap();
        assert_eq!(frame, Frame::Simple("OK".into()));
        assert_eq!(consumed, 5);
    }

    #[test]
    fn deeply_nested_array_rejected() {
        // build a frame nested 65 levels deep (exceeds MAX_NESTING_DEPTH of 64)
        let mut buf = Vec::new();
        for _ in 0..65 {
            buf.extend_from_slice(b"*1\r\n");
        }
        buf.extend_from_slice(b":1\r\n"); // leaf value

        let err = parse_frame(&buf).unwrap_err();
        assert!(
            matches!(err, ProtocolError::NestingTooDeep(64)),
            "expected NestingTooDeep, got {err:?}"
        );
    }

    #[test]
    fn nesting_at_limit_accepted() {
        // exactly 64 levels deep — should succeed
        let mut buf = Vec::new();
        for _ in 0..64 {
            buf.extend_from_slice(b"*1\r\n");
        }
        buf.extend_from_slice(b":1\r\n");

        let result = parse_frame(&buf);
        assert!(result.is_ok(), "64 levels of nesting should be accepted");
        assert!(result.unwrap().is_some());
    }
}