redis_protocol/resp2/
encode.rs

1//! Functions for encoding Frames into the RESP2 protocol.
2//!
3//! <https://redis.io/topics/protocol#resp-protocol-description>
4
5use crate::resp2::types::*;
6use crate::resp2::utils::{self as resp2_utils};
7use crate::types::{RedisProtocolError, CRLF};
8use crate::utils;
9use crate::alloc::string::ToString;
10use alloc::vec::Vec;
11use bytes::BytesMut;
12use cookie_factory::GenError;
13
14fn gen_simplestring<'a>(x: (&'a mut [u8], usize), data: &[u8]) -> Result<(&'a mut [u8], usize), GenError> {
15  encode_checks!(x, resp2_utils::simplestring_encode_len(data));
16
17  do_gen!(
18    x,
19    gen_be_u8!(FrameKind::SimpleString.to_byte()) >> gen_slice!(data) >> gen_slice!(CRLF.as_bytes())
20  )
21}
22
23fn gen_error<'a>(x: (&'a mut [u8], usize), data: &str) -> Result<(&'a mut [u8], usize), GenError> {
24  encode_checks!(x, resp2_utils::error_encode_len(data));
25
26  do_gen!(
27    x,
28    gen_be_u8!(FrameKind::Error.to_byte()) >> gen_slice!(data.as_bytes()) >> gen_slice!(CRLF.as_bytes())
29  )
30}
31
32fn gen_integer<'a>(x: (&'a mut [u8], usize), data: &i64) -> Result<(&'a mut [u8], usize), GenError> {
33  encode_checks!(x, resp2_utils::integer_encode_len(data));
34
35  do_gen!(
36    x,
37    gen_be_u8!(FrameKind::Integer.to_byte()) >> gen_slice!(data.to_string().as_bytes()) >> gen_slice!(CRLF.as_bytes())
38  )
39}
40
41fn gen_bulkstring<'a>(x: (&'a mut [u8], usize), data: &[u8]) -> Result<(&'a mut [u8], usize), GenError> {
42  encode_checks!(x, resp2_utils::bulkstring_encode_len(data));
43
44  do_gen!(
45    x,
46    gen_be_u8!(FrameKind::BulkString.to_byte())
47      >> gen_slice!(data.len().to_string().as_bytes())
48      >> gen_slice!(CRLF.as_bytes())
49      >> gen_slice!(data)
50      >> gen_slice!(CRLF.as_bytes())
51  )
52}
53
54fn gen_null(x: (&mut [u8], usize)) -> Result<(&mut [u8], usize), GenError> {
55  encode_checks!(x, NULL.as_bytes().len());
56
57  do_gen!(x, gen_slice!(NULL.as_bytes()))
58}
59
60fn gen_array<'a>(x: (&'a mut [u8], usize), data: &Vec<Frame>) -> Result<(&'a mut [u8], usize), GenError> {
61  encode_checks!(x, resp2_utils::array_encode_len(data)?);
62
63  let mut x = do_gen!(
64    x,
65    gen_be_u8!(FrameKind::Array.to_byte())
66      >> gen_slice!(data.len().to_string().as_bytes())
67      >> gen_slice!(CRLF.as_bytes())
68  )?;
69
70  for frame in data.iter() {
71    x = match frame {
72      Frame::SimpleString(ref s) => gen_simplestring(x, &s)?,
73      Frame::BulkString(ref b) => gen_bulkstring(x, &b)?,
74      Frame::Null => gen_null(x)?,
75      Frame::Error(ref s) => gen_error(x, s)?,
76      Frame::Array(ref frames) => gen_array(x, frames)?,
77      Frame::Integer(ref i) => gen_integer(x, i)?,
78    };
79  }
80
81  // no trailing CRLF here, the inner values add that
82  Ok(x)
83}
84
85fn attempt_encoding(buf: &mut [u8], offset: usize, frame: &Frame) -> Result<usize, GenError> {
86  match *frame {
87    Frame::BulkString(ref b) => gen_bulkstring((buf, offset), b).map(|(_, l)| l),
88    Frame::Null => gen_null((buf, offset)).map(|(_, l)| l),
89    Frame::Array(ref frames) => gen_array((buf, offset), frames).map(|(_, l)| l),
90    Frame::Error(ref s) => gen_error((buf, offset), s).map(|(_, l)| l),
91    Frame::SimpleString(ref s) => gen_simplestring((buf, offset), s).map(|(_, l)| l),
92    Frame::Integer(ref i) => gen_integer((buf, offset), i).map(|(_, l)| l),
93  }
94}
95
96/// Attempt to encode a frame into `buf`, assuming a starting offset of 0.
97///
98/// The caller is responsible for extending the buffer if a `RedisProtocolErrorKind::BufferTooSmall` is returned.
99pub fn encode(buf: &mut [u8], offset: usize, frame: &Frame) -> Result<usize, RedisProtocolError> {
100  attempt_encoding(buf, offset, frame).map_err(|e| e.into())
101}
102
103/// Attempt to encode a frame into `buf`, extending the buffer as needed.
104///
105/// Returns the number of bytes encoded.
106pub fn encode_bytes(buf: &mut BytesMut, frame: &Frame) -> Result<usize, RedisProtocolError> {
107  let offset = buf.len();
108
109  loop {
110    match attempt_encoding(buf, offset, frame) {
111      Ok(size) => return Ok(size),
112      Err(e) => match e {
113        GenError::BufferTooSmall(amt) => utils::zero_extend(buf, amt),
114        _ => return Err(e.into()),
115      },
116    }
117  }
118}
119
120#[cfg(test)]
121mod tests {
122  use super::*;
123  use crate::utils::*;
124
125  const PADDING: &'static str = "foobar";
126
127  fn encode_and_verify_empty(input: &Frame, expected: &str) {
128    let mut buf = BytesMut::new();
129
130    let len = match encode_bytes(&mut buf, input) {
131      Ok(l) => l,
132      Err(e) => panic!("{:?}", e),
133    };
134
135    assert_eq!(buf, expected.as_bytes(), "empty buf contents match");
136    assert_eq!(len, expected.as_bytes().len(), "empty expected len is correct");
137  }
138
139  fn encode_and_verify_non_empty(input: &Frame, expected: &str) {
140    let mut buf = BytesMut::new();
141    buf.extend_from_slice(PADDING.as_bytes());
142
143    let len = match encode_bytes(&mut buf, input) {
144      Ok(l) => l,
145      Err(e) => panic!("{:?}", e),
146    };
147    let padded = vec![PADDING, expected].join("");
148
149    assert_eq!(buf, padded.as_bytes(), "padded buf contents match");
150    assert_eq!(len, padded.as_bytes().len(), "padded expected len is correct");
151  }
152
153  fn encode_raw_and_verify_empty(input: &Frame, expected: &str) {
154    let mut buf = Vec::from(&ZEROED_KB[0..expected.as_bytes().len()]);
155
156    let len = match encode(&mut buf, 0, input) {
157      Ok(l) => l,
158      Err(e) => panic!("{:?}", e),
159    };
160
161    assert_eq!(buf, expected.as_bytes(), "empty buf contents match");
162    assert_eq!(len, expected.as_bytes().len(), "empty expected len is correct");
163  }
164
165  #[test]
166  fn should_encode_llen_req_example() {
167    let expected = "*2\r\n$4\r\nLLEN\r\n$6\r\nmylist\r\n";
168    let input = Frame::Array(vec![
169      Frame::BulkString("LLEN".into()),
170      Frame::BulkString("mylist".into()),
171    ]);
172
173    encode_and_verify_empty(&input, expected);
174    encode_and_verify_non_empty(&input, expected);
175  }
176
177  #[test]
178  fn should_encode_incr_req_example() {
179    let expected = "*2\r\n$4\r\nINCR\r\n$5\r\nmykey\r\n";
180    let input = Frame::Array(vec![
181      Frame::BulkString("INCR".into()),
182      Frame::BulkString("mykey".into()),
183    ]);
184
185    encode_and_verify_empty(&input, expected);
186    encode_and_verify_non_empty(&input, expected);
187  }
188
189  #[test]
190  fn should_encode_bitcount_req_example() {
191    let expected = "*2\r\n$8\r\nBITCOUNT\r\n$5\r\nmykey\r\n";
192    let input = Frame::Array(vec![
193      Frame::BulkString("BITCOUNT".into()),
194      Frame::BulkString("mykey".into()),
195    ]);
196
197    encode_and_verify_empty(&input, expected);
198    encode_and_verify_non_empty(&input, expected);
199  }
200
201  #[test]
202  fn should_encode_array_bulk_string_test() {
203    let expected = "*3\r\n$5\r\nWATCH\r\n$6\r\nWIBBLE\r\n$9\r\nfooBARbaz\r\n";
204    let input = Frame::Array(vec![
205      Frame::BulkString("WATCH".into()),
206      Frame::BulkString("WIBBLE".into()),
207      Frame::BulkString("fooBARbaz".into()),
208    ]);
209
210    encode_and_verify_empty(&input, expected);
211    encode_and_verify_non_empty(&input, expected);
212  }
213
214  #[test]
215  fn should_encode_array_null_test() {
216    let expected = "*3\r\n$4\r\nHSET\r\n$3\r\nfoo\r\n$-1\r\n";
217    let input = Frame::Array(vec![
218      Frame::BulkString("HSET".into()),
219      Frame::BulkString("foo".into()),
220      Frame::Null,
221    ]);
222
223    encode_and_verify_empty(&input, expected);
224    encode_and_verify_non_empty(&input, expected);
225  }
226
227  #[test]
228  fn should_encode_raw_llen_req_example() {
229    let expected = "*2\r\n$4\r\nLLEN\r\n$6\r\nmylist\r\n";
230    let input = Frame::Array(vec![
231      Frame::BulkString("LLEN".into()),
232      Frame::BulkString("mylist".into()),
233    ]);
234
235    encode_raw_and_verify_empty(&input, expected);
236  }
237
238  #[test]
239  fn should_encode_raw_incr_req_example() {
240    let expected = "*2\r\n$4\r\nINCR\r\n$5\r\nmykey\r\n";
241    let input = Frame::Array(vec![
242      Frame::BulkString("INCR".into()),
243      Frame::BulkString("mykey".into()),
244    ]);
245
246    encode_raw_and_verify_empty(&input, expected);
247  }
248
249  #[test]
250  fn should_encode_raw_bitcount_req_example() {
251    let expected = "*2\r\n$8\r\nBITCOUNT\r\n$5\r\nmykey\r\n";
252    let input = Frame::Array(vec![
253      Frame::BulkString("BITCOUNT".into()),
254      Frame::BulkString("mykey".into()),
255    ]);
256
257    encode_raw_and_verify_empty(&input, expected);
258  }
259
260  #[test]
261  fn should_encode_raw_array_bulk_string_test() {
262    let expected = "*3\r\n$5\r\nWATCH\r\n$6\r\nWIBBLE\r\n$9\r\nfooBARbaz\r\n";
263    let input = Frame::Array(vec![
264      Frame::BulkString("WATCH".into()),
265      Frame::BulkString("WIBBLE".into()),
266      Frame::BulkString("fooBARbaz".into()),
267    ]);
268
269    encode_raw_and_verify_empty(&input, expected);
270  }
271
272  #[test]
273  fn should_encode_raw_array_null_test() {
274    let expected = "*3\r\n$4\r\nHSET\r\n$3\r\nfoo\r\n$-1\r\n";
275    let input = Frame::Array(vec![
276      Frame::BulkString("HSET".into()),
277      Frame::BulkString("foo".into()),
278      Frame::Null,
279    ]);
280
281    encode_raw_and_verify_empty(&input, expected);
282  }
283
284  #[test]
285  fn should_encode_moved_error() {
286    let expected = "-MOVED 3999 127.0.0.1:6381\r\n";
287    let input = Frame::Error("MOVED 3999 127.0.0.1:6381".into());
288
289    encode_and_verify_empty(&input, expected);
290    encode_and_verify_non_empty(&input, expected);
291  }
292
293  #[test]
294  fn should_encode_ask_error() {
295    let expected = "-ASK 3999 127.0.0.1:6381\r\n";
296    let input = Frame::Error("ASK 3999 127.0.0.1:6381".into());
297
298    encode_and_verify_empty(&input, expected);
299    encode_and_verify_non_empty(&input, expected);
300  }
301
302  #[test]
303  fn should_encode_error() {
304    let expected = "-WRONGTYPE Operation against a key holding the wrong kind of value\r\n";
305    let input = Frame::Error("WRONGTYPE Operation against a key holding the wrong kind of value".into());
306
307    encode_and_verify_empty(&input, expected);
308    encode_and_verify_non_empty(&input, expected);
309  }
310
311  #[test]
312  fn should_encode_simplestring() {
313    let expected = "+OK\r\n";
314    let input = Frame::SimpleString("OK".into());
315
316    encode_and_verify_empty(&input, expected);
317    encode_and_verify_non_empty(&input, expected);
318  }
319
320  #[test]
321  fn should_encode_integer() {
322    let i1_expected = ":1000\r\n";
323    let i1_input = Frame::Integer(1000);
324
325    encode_and_verify_empty(&i1_input, i1_expected);
326    encode_and_verify_non_empty(&i1_input, i1_expected);
327  }
328
329  #[test]
330  fn should_encode_negative_integer() {
331    let i2_expected = ":-1000\r\n";
332    let i2_input = Frame::Integer(-1000);
333
334    encode_and_verify_empty(&i2_input, i2_expected);
335    encode_and_verify_non_empty(&i2_input, i2_expected);
336  }
337}
redis_protocol/resp2/encode.rs

redis_protocol/resp2/
encode.rs
