redis_protocol/resp2/
utils.rs

1use crate::{
2  error::RedisProtocolError,
3  resp2::types::{OwnedFrame, RangeFrame},
4  utils::digits_in_usize,
5};
6use alloc::{borrow::ToOwned, string::String, vec::Vec};
7
8#[cfg(feature = "bytes")]
9use crate::resp2::types::BytesFrame;
10#[cfg(feature = "bytes")]
11use bytes::{Bytes, BytesMut};
12#[cfg(feature = "bytes")]
13use bytes_utils::Str;
14
15pub fn bulkstring_encode_len(b: &[u8]) -> usize {
16  1 + digits_in_usize(b.len()) + 2 + b.len() + 2
17}
18
19pub fn simplestring_encode_len(s: &[u8]) -> usize {
20  1 + s.len() + 2
21}
22
23pub fn error_encode_len(s: &str) -> usize {
24  1 + s.as_bytes().len() + 2
25}
26
27pub fn integer_encode_len(i: i64, int_as_bulkstring: bool) -> usize {
28  let prefix = if i < 0 { 1 } else { 0 };
29  let digits = digits_in_usize(i.unsigned_abs() as usize);
30
31  if int_as_bulkstring {
32    1 + digits_in_usize(digits + prefix) + 2 + prefix + digits + 2
33  } else {
34    1 + digits + 2 + prefix
35  }
36}
37
38/// Move or copy the contents of `buf` based on the ranges in the provided frame.
39pub fn build_owned_frame(buf: &[u8], frame: &RangeFrame) -> Result<OwnedFrame, RedisProtocolError> {
40  Ok(match frame {
41    RangeFrame::Error((start, end)) => OwnedFrame::Error(String::from_utf8(buf[*start .. *end].to_vec())?),
42    RangeFrame::SimpleString((start, end)) => OwnedFrame::SimpleString(buf[*start .. *end].to_owned()),
43    RangeFrame::BulkString((start, end)) => OwnedFrame::BulkString(buf[*start .. *end].to_owned()),
44    RangeFrame::Integer(i) => OwnedFrame::Integer(*i),
45    RangeFrame::Null => OwnedFrame::Null,
46    RangeFrame::Array(frames) => {
47      // FIXME is there a cleaner way to do this with results?
48      let mut out = Vec::with_capacity(frames.len());
49      for frame in frames.iter() {
50        out.push(build_owned_frame(buf, frame)?);
51      }
52      OwnedFrame::Array(out)
53    },
54  })
55}
56
57/// Use the `Bytes` interface to create owned views into the provided buffer for each of the provided range frames.
58#[cfg(feature = "bytes")]
59pub fn build_bytes_frame(buf: &Bytes, frame: &RangeFrame) -> Result<BytesFrame, RedisProtocolError> {
60  Ok(match frame {
61    RangeFrame::Error((start, end)) => {
62      let bytes = buf.slice(start .. end);
63      BytesFrame::Error(Str::from_inner(bytes)?)
64    },
65    RangeFrame::SimpleString((start, end)) => {
66      let bytes = buf.slice(start .. end);
67      BytesFrame::SimpleString(bytes)
68    },
69    RangeFrame::BulkString((start, end)) => {
70      let bytes = buf.slice(start .. end);
71      BytesFrame::BulkString(bytes)
72    },
73    RangeFrame::Integer(i) => BytesFrame::Integer(*i),
74    RangeFrame::Null => BytesFrame::Null,
75    RangeFrame::Array(frames) => {
76      // FIXME is there a cleaner way to do this with results?
77      let mut out = Vec::with_capacity(frames.len());
78      for frame in frames.iter() {
79        out.push(build_bytes_frame(buf, frame)?);
80      }
81      BytesFrame::Array(out)
82    },
83  })
84}
85
86/// Split off and [freeze](bytes::BytesMut::freeze) `amt` bytes from `buf` and return owned views into the buffer
87/// based on the provided range frame.
88///
89/// The returned `Bytes` represents the `Bytes` buffer that was sliced off `buf`. The returned frames hold
90/// owned `Bytes` views to slices within this buffer.
91#[cfg(feature = "bytes")]
92pub fn freeze_parse(
93  buf: &mut BytesMut,
94  frame: &RangeFrame,
95  amt: usize,
96) -> Result<(BytesFrame, Bytes), RedisProtocolError> {
97  let buffer = buf.split_to(amt).freeze();
98  let frame = build_bytes_frame(&buffer, frame)?;
99  Ok((frame, buffer))
100}
101
102#[cfg(test)]
103mod tests {
104  use super::*;
105
106  #[test]
107  fn should_get_encode_len_simplestring() {
108    let ss1 = "Ok";
109    let ss2 = "FooBarBaz";
110    let ss3 = "-&#$@9232";
111
112    assert_eq!(simplestring_encode_len(ss1.as_bytes()), 5);
113    assert_eq!(simplestring_encode_len(ss2.as_bytes()), 12);
114    assert_eq!(simplestring_encode_len(ss3.as_bytes()), 12);
115  }
116
117  #[test]
118  fn should_get_encode_len_error() {
119    let e1 = "MOVED 3999 127.0.0.1:6381";
120    let e2 = "ERR unknown command 'foobar'";
121    let e3 = "WRONGTYPE Operation against a key holding the wrong kind of value";
122
123    assert_eq!(error_encode_len(e1), 28);
124    assert_eq!(error_encode_len(e2), 31);
125    assert_eq!(error_encode_len(e3), 68);
126  }
127
128  #[test]
129  fn should_get_encode_len_integer() {
130    assert_eq!(integer_encode_len(38473, false), 8);
131    assert_eq!(integer_encode_len(-74834, false), 9);
132  }
133
134  #[test]
135  fn should_get_encode_len_integer_as_bulkstring() {
136    // $5\r\n38473\r\n
137    assert_eq!(integer_encode_len(38473, true), 11);
138    // $6\r\n-74834\r\n
139    assert_eq!(integer_encode_len(-74834, true), 12);
140  }
141}
redis_protocol/resp2/utils.rs

redis_protocol/resp2/
utils.rs
