Skip to main content

ocular_protocol/
resp.rs

1use anyhow::{Result, bail};
2
3/// Redis RESP value
4#[derive(Debug, Clone, PartialEq)]
5pub enum RespValue {
6    Simple(String),
7    Error(String),
8    Integer(i64),
9    Bulk(Option<Vec<u8>>),
10    Array(Option<Vec<RespValue>>),
11}
12
13impl RespValue {
14    /// Format a RESP command as a readable string, e.g. "SET key value"
15    pub fn to_command_string(&self) -> String {
16        match self {
17            RespValue::Array(Some(parts)) => {
18                parts.iter().map(|p| match p {
19                    RespValue::Bulk(Some(b)) => String::from_utf8_lossy(b).to_string(),
20                    other => format!("{:?}", other),
21                }).collect::<Vec<_>>().join(" ")
22            }
23            RespValue::Simple(s) => s.clone(),
24            RespValue::Error(e) => format!("ERR: {}", e),
25            RespValue::Integer(i) => i.to_string(),
26            RespValue::Bulk(Some(b)) => String::from_utf8_lossy(b).to_string(),
27            _ => String::from("(nil)"),
28        }
29    }
30}
31
32/// Parse a complete RESP value from a byte buffer. Returns (value, bytes consumed).
33pub fn parse_resp(buf: &[u8]) -> Result<Option<(RespValue, usize)>> {
34    if buf.is_empty() {
35        return Ok(None);
36    }
37    parse_value(buf, 0)
38}
39
40fn parse_value(buf: &[u8], pos: usize) -> Result<Option<(RespValue, usize)>> {
41    if pos >= buf.len() {
42        return Ok(None);
43    }
44    match buf[pos] {
45        b'+' => parse_simple(buf, pos),
46        b'-' => parse_error(buf, pos),
47        b':' => parse_integer(buf, pos),
48        b'$' => parse_bulk(buf, pos),
49        b'*' => parse_array(buf, pos),
50        _ => bail!("unknown RESP type byte: {:02x}", buf[pos]),
51    }
52}
53
54fn find_crlf(buf: &[u8], start: usize) -> Option<usize> {
55    buf[start..].windows(2).position(|w| w == b"\r\n").map(|i| start + i)
56}
57
58fn parse_line(buf: &[u8], pos: usize) -> Option<(&[u8], usize)> {
59    find_crlf(buf, pos).map(|end| (&buf[pos..end], end + 2))
60}
61
62fn parse_simple(buf: &[u8], pos: usize) -> Result<Option<(RespValue, usize)>> {
63    match parse_line(buf, pos + 1) {
64        Some((line, next)) => Ok(Some((RespValue::Simple(String::from_utf8_lossy(line).to_string()), next))),
65        None => Ok(None),
66    }
67}
68
69fn parse_error(buf: &[u8], pos: usize) -> Result<Option<(RespValue, usize)>> {
70    match parse_line(buf, pos + 1) {
71        Some((line, next)) => Ok(Some((RespValue::Error(String::from_utf8_lossy(line).to_string()), next))),
72        None => Ok(None),
73    }
74}
75
76fn parse_integer(buf: &[u8], pos: usize) -> Result<Option<(RespValue, usize)>> {
77    match parse_line(buf, pos + 1) {
78        Some((line, next)) => {
79            let s = std::str::from_utf8(line)?;
80            Ok(Some((RespValue::Integer(s.parse()?), next)))
81        }
82        None => Ok(None),
83    }
84}
85
86fn parse_bulk(buf: &[u8], pos: usize) -> Result<Option<(RespValue, usize)>> {
87    let Some((line, next)) = parse_line(buf, pos + 1) else { return Ok(None) };
88    let len: i64 = std::str::from_utf8(line)?.parse()?;
89    if len < 0 {
90        return Ok(Some((RespValue::Bulk(None), next)));
91    }
92    let len = len as usize;
93    let end = next + len + 2; // data + \r\n
94    if buf.len() < end {
95        return Ok(None);
96    }
97    Ok(Some((RespValue::Bulk(Some(buf[next..next + len].to_vec())), end)))
98}
99
100fn parse_array(buf: &[u8], pos: usize) -> Result<Option<(RespValue, usize)>> {
101    let Some((line, mut next)) = parse_line(buf, pos + 1) else { return Ok(None) };
102    let count: i64 = std::str::from_utf8(line)?.parse()?;
103    if count < 0 {
104        return Ok(Some((RespValue::Array(None), next)));
105    }
106    let mut items = Vec::with_capacity(count as usize);
107    for _ in 0..count {
108        match parse_value(buf, next)? {
109            Some((val, consumed)) => {
110                items.push(val);
111                next = consumed;
112            }
113            None => return Ok(None),
114        }
115    }
116    Ok(Some((RespValue::Array(Some(items)), next)))
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122
123    #[test]
124    fn test_parse_simple() {
125        let input = b"+OK\r\n";
126        let (val, n) = parse_resp(input).unwrap().unwrap();
127        assert_eq!(val, RespValue::Simple("OK".into()));
128        assert_eq!(n, 5);
129    }
130
131    #[test]
132    fn test_parse_array_command() {
133        // *3\r\n$3\r\nSET\r\n$3\r\nkey\r\n$5\r\nvalue\r\n
134        let input = b"*3\r\n$3\r\nSET\r\n$3\r\nkey\r\n$5\r\nvalue\r\n";
135        let (val, _) = parse_resp(input).unwrap().unwrap();
136        assert_eq!(val.to_command_string(), "SET key value");
137    }
138}