Skip to main content

feox_server/protocol/
resp.rs

1use bytes::{Bytes, BytesMut};
2use memchr::memchr2;
3use std::str;
4
5/// RESP (REdis Serialization Protocol) parser
6pub struct RespParser {
7    buffer: BytesMut,
8    position: usize,
9}
10
11#[derive(Debug, Clone)]
12pub enum RespValue {
13    SimpleString(Bytes),
14    Error(String),
15    Integer(i64),
16    BulkString(Option<Bytes>),
17    Array(Option<Vec<RespValue>>),
18}
19
20impl RespParser {
21    pub fn new() -> Self {
22        Self {
23            buffer: BytesMut::with_capacity(16 * 1024),
24            position: 0,
25        }
26    }
27
28    /// Feed data into the parser
29    pub fn feed(&mut self, data: &[u8]) {
30        self.buffer.extend_from_slice(data);
31    }
32
33    /// Parse next complete RESP value
34    pub fn parse_next(&mut self) -> Result<Option<RespValue>, String> {
35        if self.position >= self.buffer.len() {
36            return Ok(None);
37        }
38
39        let remaining = &self.buffer[self.position..];
40
41        match self.parse_value(remaining) {
42            Ok(Some((value, consumed))) => {
43                self.position += consumed;
44
45                // Compact buffer if needed
46                if self.position > self.buffer.len() / 2 {
47                    let _ = self.buffer.split_to(self.position);
48                    self.position = 0;
49                }
50
51                Ok(Some(value))
52            }
53            Ok(None) => Ok(None), // Need more data
54            Err(e) => Err(e),
55        }
56    }
57
58    /// Parse a RESP value from buffer
59    fn parse_value(&self, buf: &[u8]) -> Result<Option<(RespValue, usize)>, String> {
60        if buf.is_empty() {
61            return Ok(None);
62        }
63
64        match buf[0] {
65            b'+' => self.parse_simple_string(buf),
66            b'-' => self.parse_error(buf),
67            b':' => self.parse_integer(buf),
68            b'$' => self.parse_bulk_string(buf),
69            b'*' => self.parse_array(buf),
70            _ => Err(format!("Invalid RESP type: {}", buf[0] as char)),
71        }
72    }
73
74    /// Parse simple string: +OK\r\n
75    fn parse_simple_string(&self, buf: &[u8]) -> Result<Option<(RespValue, usize)>, String> {
76        if let Some(end) = find_crlf(buf) {
77            let data = Bytes::copy_from_slice(&buf[1..end]);
78            Ok(Some((RespValue::SimpleString(data), end + 2)))
79        } else {
80            Ok(None)
81        }
82    }
83
84    /// Parse error: -ERR message\r\n
85    fn parse_error(&self, buf: &[u8]) -> Result<Option<(RespValue, usize)>, String> {
86        if let Some(end) = find_crlf(buf) {
87            let msg = str::from_utf8(&buf[1..end])
88                .map_err(|_| "Invalid UTF-8 in error")?
89                .to_string();
90            Ok(Some((RespValue::Error(msg), end + 2)))
91        } else {
92            Ok(None)
93        }
94    }
95
96    /// Parse integer: :123\r\n
97    fn parse_integer(&self, buf: &[u8]) -> Result<Option<(RespValue, usize)>, String> {
98        if let Some(end) = find_crlf(buf) {
99            let num_str = str::from_utf8(&buf[1..end]).map_err(|_| "Invalid UTF-8 in integer")?;
100            let num = num_str
101                .parse::<i64>()
102                .map_err(|_| format!("Invalid integer: {}", num_str))?;
103            Ok(Some((RespValue::Integer(num), end + 2)))
104        } else {
105            Ok(None)
106        }
107    }
108
109    /// Parse bulk string: $6\r\nfoobar\r\n or $-1\r\n (null)
110    fn parse_bulk_string(&self, buf: &[u8]) -> Result<Option<(RespValue, usize)>, String> {
111        // Find length line
112        let len_end = match find_crlf(buf) {
113            Some(pos) => pos,
114            None => return Ok(None),
115        };
116
117        // Parse length
118        let len_str =
119            str::from_utf8(&buf[1..len_end]).map_err(|_| "Invalid UTF-8 in bulk string length")?;
120        let len = len_str
121            .parse::<i64>()
122            .map_err(|_| format!("Invalid bulk string length: {}", len_str))?;
123
124        if len < 0 {
125            // Null bulk string
126            return Ok(Some((RespValue::BulkString(None), len_end + 2)));
127        }
128
129        let len = len as usize;
130        let data_start = len_end + 2;
131        let data_end = data_start + len;
132
133        // Check if we have enough data
134        if buf.len() < data_end + 2 {
135            return Ok(None);
136        }
137
138        // Verify CRLF after data
139        if buf[data_end] != b'\r' || buf[data_end + 1] != b'\n' {
140            return Err("Missing CRLF after bulk string".to_string());
141        }
142
143        // Zero-copy: create Bytes view into buffer
144        let data = Bytes::copy_from_slice(&buf[data_start..data_end]);
145        Ok(Some((RespValue::BulkString(Some(data)), data_end + 2)))
146    }
147
148    /// Parse array: *2\r\n$3\r\nGET\r\n$3\r\nkey\r\n
149    fn parse_array(&self, buf: &[u8]) -> Result<Option<(RespValue, usize)>, String> {
150        // Find length line
151        let len_end = match find_crlf(buf) {
152            Some(pos) => pos,
153            None => return Ok(None),
154        };
155
156        // Parse length
157        let len_str =
158            str::from_utf8(&buf[1..len_end]).map_err(|_| "Invalid UTF-8 in array length")?;
159        let len = len_str
160            .parse::<i64>()
161            .map_err(|_| format!("Invalid array length: {}", len_str))?;
162
163        if len < 0 {
164            // Null array
165            return Ok(Some((RespValue::Array(None), len_end + 2)));
166        }
167
168        let len = len as usize;
169        let mut elements = Vec::with_capacity(len);
170        let mut pos = len_end + 2;
171
172        // Parse array elements
173        for _ in 0..len {
174            match self.parse_value(&buf[pos..])? {
175                Some((value, consumed)) => {
176                    elements.push(value);
177                    pos += consumed;
178                }
179                None => return Ok(None), // Need more data
180            }
181        }
182
183        Ok(Some((RespValue::Array(Some(elements)), pos)))
184    }
185}
186
187/// Find CRLF in buffer
188#[inline]
189fn find_crlf(buf: &[u8]) -> Option<usize> {
190    let mut pos = 0;
191    while pos < buf.len() - 1 {
192        if let Some(cr_pos) = memchr2(b'\r', b'\n', &buf[pos..]) {
193            let cr_pos = pos + cr_pos;
194            if cr_pos < buf.len() - 1 && buf[cr_pos] == b'\r' && buf[cr_pos + 1] == b'\n' {
195                return Some(cr_pos);
196            }
197            pos = cr_pos + 1;
198        } else {
199            break;
200        }
201    }
202    None
203}
204
205/// Format response as RESP with pre-allocated buffer
206pub fn format_resp_response(value: &RespValue) -> Vec<u8> {
207    let mut result = Vec::with_capacity(estimate_resp_size(value));
208    write_resp_value(&mut result, value);
209    result
210}
211
212/// Write RESP value directly to buffer (zero-copy when possible)
213pub fn write_resp_value(buf: &mut Vec<u8>, value: &RespValue) {
214    match value {
215        RespValue::SimpleString(s) => {
216            buf.push(b'+');
217            buf.extend_from_slice(s);
218            buf.extend_from_slice(b"\r\n");
219        }
220        RespValue::Error(e) => {
221            buf.push(b'-');
222            buf.extend_from_slice(e.as_bytes());
223            buf.extend_from_slice(b"\r\n");
224        }
225        RespValue::Integer(n) => {
226            buf.push(b':');
227            let mut num_buf = itoa::Buffer::new();
228            buf.extend_from_slice(num_buf.format(*n).as_bytes());
229            buf.extend_from_slice(b"\r\n");
230        }
231        RespValue::BulkString(Some(s)) => {
232            buf.push(b'$');
233            let mut num_buf = itoa::Buffer::new();
234            buf.extend_from_slice(num_buf.format(s.len()).as_bytes());
235            buf.extend_from_slice(b"\r\n");
236            buf.extend_from_slice(s);
237            buf.extend_from_slice(b"\r\n");
238        }
239        RespValue::BulkString(None) => {
240            buf.extend_from_slice(b"$-1\r\n");
241        }
242        RespValue::Array(Some(arr)) => {
243            buf.push(b'*');
244            let mut num_buf = itoa::Buffer::new();
245            buf.extend_from_slice(num_buf.format(arr.len()).as_bytes());
246            buf.extend_from_slice(b"\r\n");
247            for item in arr {
248                write_resp_value(buf, item);
249            }
250        }
251        RespValue::Array(None) => {
252            buf.extend_from_slice(b"*-1\r\n");
253        }
254    }
255}
256
257/// Estimate the size needed for a RESP value (for pre-allocation)
258#[inline]
259fn estimate_resp_size(value: &RespValue) -> usize {
260    match value {
261        RespValue::SimpleString(s) => s.len() + 3,
262        RespValue::Error(e) => e.len() + 3,
263        RespValue::Integer(_) => 24, // :number\r\n
264        RespValue::BulkString(Some(s)) => s.len() + 20,
265        RespValue::BulkString(None) => 5,
266        RespValue::Array(Some(arr)) => {
267            let mut size = 10; // array header
268            for item in arr {
269                size += estimate_resp_size(item);
270            }
271            size
272        }
273        RespValue::Array(None) => 5,
274    }
275}
276
277impl Default for RespParser {
278    fn default() -> Self {
279        Self::new()
280    }
281}