rust_cache_server/
protocol.rs

1use std::fmt;
2
3use bytes::{Buf, BytesMut};
4use thiserror::Error;
5
6#[derive(Debug, Error)]
7pub enum RespError {
8    #[error("incomplete frame")]
9    Incomplete,
10    #[error("protocol error: {0}")]
11    Protocol(String),
12}
13
14#[derive(Debug, Clone)]
15pub enum RespFrame {
16    SimpleString(String),
17    Error(String),
18    Integer(i64),
19    BulkString(Option<Vec<u8>>),
20    Array(Option<Vec<RespFrame>>),
21}
22
23impl fmt::Display for RespFrame {
24    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25        match self {
26            RespFrame::SimpleString(s) => write!(f, "{}", s),
27            RespFrame::Error(s) => write!(f, "ERR {}", s),
28            RespFrame::Integer(i) => write!(f, "{}", i),
29            RespFrame::BulkString(Some(b)) => write!(f, "{}", String::from_utf8_lossy(b)),
30            RespFrame::BulkString(None) => write!(f, "(nil)"),
31            RespFrame::Array(Some(arr)) => {
32                write!(f, "[")?;
33                for (i, a) in arr.iter().enumerate() {
34                    if i > 0 {
35                        write!(f, ", ")?;
36                    }
37                    write!(f, "{a}")?;
38                }
39                write!(f, "]")
40            }
41            RespFrame::Array(None) => write!(f, "(nil array)"),
42        }
43    }
44}
45
46fn read_line(buf: &mut BytesMut) -> Result<Option<Vec<u8>>, RespError> {
47    if let Some(pos) = buf.windows(2).position(|w| w == b"\r\n") {
48        let mut line = buf.split_to(pos + 2);
49        line.truncate(pos);
50        Ok(Some(line.to_vec()))
51    } else {
52        Ok(None)
53    }
54}
55
56fn parse_isize(line: &[u8]) -> Result<isize, RespError> {
57    let s = std::str::from_utf8(line).map_err(|e| RespError::Protocol(e.to_string()))?;
58    s.parse::<isize>()
59        .map_err(|e| RespError::Protocol(e.to_string()))
60}
61
62pub fn parse_frame(buf: &mut BytesMut) -> Result<Option<RespFrame>, RespError> {
63    if !buf.has_remaining() {
64        return Ok(None);
65    }
66
67    match buf[0] {
68        b'+' | b'-' | b':' => {
69            let line = read_line(buf)?.ok_or(RespError::Incomplete)?;
70            let (ty, data) = line
71                .split_first()
72                .ok_or_else(|| RespError::Protocol("empty line".into()))?;
73            let s =
74                std::str::from_utf8(data).map_err(|e| RespError::Protocol(e.to_string()))?;
75            let frame = match ty {
76                b'+' => RespFrame::SimpleString(s.to_string()),
77                b'-' => RespFrame::Error(s.to_string()),
78                b':' => RespFrame::Integer(
79                    s.parse::<i64>()
80                        .map_err(|e| RespError::Protocol(e.to_string()))?,
81                ),
82                _ => return Err(RespError::Protocol("invalid prefix".into())),
83            };
84            Ok(Some(frame))
85        }
86        b'$' => {
87            let header = read_line(buf)?.ok_or(RespError::Incomplete)?;
88            let len = parse_isize(&header[1..])?;
89            if len < 0 {
90                return Ok(Some(RespFrame::BulkString(None)));
91            }
92            let len = len as usize;
93            if buf.len() < len + 2 {
94                return Err(RespError::Incomplete);
95            }
96            let data = buf.split_to(len).to_vec();
97            if buf.len() < 2 || &buf[..2] != b"\r\n" {
98                return Err(RespError::Protocol("bulkstring missing CRLF".into()));
99            }
100            buf.advance(2);
101            Ok(Some(RespFrame::BulkString(Some(data))))
102        }
103        b'*' => {
104            let header = read_line(buf)?.ok_or(RespError::Incomplete)?;
105            let len = parse_isize(&header[1..])?;
106            if len < 0 {
107                return Ok(Some(RespFrame::Array(None)));
108            }
109            let len = len as usize;
110            let mut items = Vec::with_capacity(len);
111            for _ in 0..len {
112                let frame = parse_frame(buf)?.ok_or(RespError::Incomplete)?;
113                items.push(frame);
114            }
115            Ok(Some(RespFrame::Array(Some(items))))
116        }
117        _ => Err(RespError::Protocol("unknown RESP type".into())),
118    }
119}
120
121pub fn encode_frame(frame: &RespFrame) -> Vec<u8> {
122    match frame {
123        RespFrame::SimpleString(s) => {
124            let mut out = Vec::with_capacity(s.len() + 3);
125            out.extend_from_slice(b"+");
126            out.extend_from_slice(s.as_bytes());
127            out.extend_from_slice(b"\r\n");
128            out
129        }
130        RespFrame::Error(s) => {
131            let mut out = Vec::with_capacity(s.len() + 3);
132            out.extend_from_slice(b"-");
133            out.extend_from_slice(s.as_bytes());
134            out.extend_from_slice(b"\r\n");
135            out
136        }
137        RespFrame::Integer(i) => {
138            let s = i.to_string();
139            let mut out = Vec::with_capacity(s.len() + 3);
140            out.extend_from_slice(b":");
141            out.extend_from_slice(s.as_bytes());
142            out.extend_from_slice(b"\r\n");
143            out
144        }
145        RespFrame::BulkString(None) => b"$-1\r\n".to_vec(),
146        RespFrame::BulkString(Some(b)) => {
147            let len = b.len().to_string();
148            let mut out = Vec::with_capacity(len.len() + b.len() + 5);
149            out.extend_from_slice(b"$");
150            out.extend_from_slice(len.as_bytes());
151            out.extend_from_slice(b"\r\n");
152            out.extend_from_slice(b);
153            out.extend_from_slice(b"\r\n");
154            out
155        }
156        RespFrame::Array(None) => b"*-1\r\n".to_vec(),
157        RespFrame::Array(Some(arr)) => {
158            let len = arr.len().to_string();
159            let mut out = Vec::new();
160            out.extend_from_slice(b"*");
161            out.extend_from_slice(len.as_bytes());
162            out.extend_from_slice(b"\r\n");
163            for f in arr {
164                out.extend_from_slice(&encode_frame(f));
165            }
166            out
167        }
168    }
169}
170
171
rust_cache_server/protocol.rs

rust_cache_server/
protocol.rs
