feox_server/protocol/
resp.rs1use bytes::{Bytes, BytesMut};
2use memchr::memchr2;
3use std::str;
4
5pub 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 pub fn feed(&mut self, data: &[u8]) {
30 self.buffer.extend_from_slice(data);
31 }
32
33 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 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), Err(e) => Err(e),
55 }
56 }
57
58 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 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 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 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 fn parse_bulk_string(&self, buf: &[u8]) -> Result<Option<(RespValue, usize)>, String> {
111 let len_end = match find_crlf(buf) {
113 Some(pos) => pos,
114 None => return Ok(None),
115 };
116
117 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 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 if buf.len() < data_end + 2 {
135 return Ok(None);
136 }
137
138 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 let data = Bytes::copy_from_slice(&buf[data_start..data_end]);
145 Ok(Some((RespValue::BulkString(Some(data)), data_end + 2)))
146 }
147
148 fn parse_array(&self, buf: &[u8]) -> Result<Option<(RespValue, usize)>, String> {
150 let len_end = match find_crlf(buf) {
152 Some(pos) => pos,
153 None => return Ok(None),
154 };
155
156 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 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 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), }
181 }
182
183 Ok(Some((RespValue::Array(Some(elements)), pos)))
184 }
185}
186
187#[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
205pub 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
212pub 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#[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, RespValue::BulkString(Some(s)) => s.len() + 20,
265 RespValue::BulkString(None) => 5,
266 RespValue::Array(Some(arr)) => {
267 let mut size = 10; 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}