Skip to main content

redis_module/
redismodule.rs

1use std::borrow::Borrow;
2use std::convert::TryFrom;
3use std::ffi::CString;
4use std::fmt::Display;
5use std::ops::Deref;
6use std::os::raw::{c_char, c_int, c_void};
7use std::ptr::NonNull;
8use std::slice;
9use std::str;
10use std::str::Utf8Error;
11use std::string::FromUtf8Error;
12use std::{fmt, ptr};
13
14use serde::de::{Error, SeqAccess};
15
16pub use crate::raw;
17pub use crate::rediserror::RedisError;
18pub use crate::redisvalue::RedisValue;
19use crate::Context;
20
21/// A short-hand type that stores a [std::result::Result] with custom
22/// type and [RedisError].
23pub type RedisResult<T = RedisValue> = Result<T, RedisError>;
24/// A [RedisResult] with [RedisValue].
25pub type RedisValueResult = RedisResult<RedisValue>;
26
27impl From<RedisValue> for RedisValueResult {
28    fn from(v: RedisValue) -> Self {
29        Ok(v)
30    }
31}
32
33impl From<RedisError> for RedisValueResult {
34    fn from(v: RedisError) -> Self {
35        Err(v)
36    }
37}
38
39pub const REDIS_OK: RedisValueResult = Ok(RedisValue::SimpleStringStatic("OK"));
40pub const TYPE_METHOD_VERSION: u64 = raw::REDISMODULE_TYPE_METHOD_VERSION as u64;
41
42pub trait NextArg {
43    fn next_arg(&mut self) -> Result<RedisString, RedisError>;
44    fn next_string(&mut self) -> Result<String, RedisError>;
45    fn next_str<'a>(&mut self) -> Result<&'a str, RedisError>;
46    fn next_i64(&mut self) -> Result<i64, RedisError>;
47    fn next_u64(&mut self) -> Result<u64, RedisError>;
48    fn next_f64(&mut self) -> Result<f64, RedisError>;
49    fn done(&mut self) -> Result<(), RedisError>;
50}
51
52impl<T> NextArg for T
53where
54    T: Iterator<Item = RedisString>,
55{
56    #[inline]
57    fn next_arg(&mut self) -> Result<RedisString, RedisError> {
58        self.next().ok_or(RedisError::WrongArity)
59    }
60
61    #[inline]
62    fn next_string(&mut self) -> Result<String, RedisError> {
63        self.next()
64            .map_or(Err(RedisError::WrongArity), |v| Ok(v.to_string_lossy()))
65    }
66
67    #[inline]
68    fn next_str<'a>(&mut self) -> Result<&'a str, RedisError> {
69        self.next()
70            .map_or(Err(RedisError::WrongArity), |v| v.try_as_str())
71    }
72
73    #[inline]
74    fn next_i64(&mut self) -> Result<i64, RedisError> {
75        self.next()
76            .map_or(Err(RedisError::WrongArity), |v| v.parse_integer())
77    }
78
79    #[inline]
80    fn next_u64(&mut self) -> Result<u64, RedisError> {
81        self.next()
82            .map_or(Err(RedisError::WrongArity), |v| v.parse_unsigned_integer())
83    }
84
85    #[inline]
86    fn next_f64(&mut self) -> Result<f64, RedisError> {
87        self.next()
88            .map_or(Err(RedisError::WrongArity), |v| v.parse_float())
89    }
90
91    /// Return an error if there are any more arguments
92    #[inline]
93    fn done(&mut self) -> Result<(), RedisError> {
94        self.next().map_or(Ok(()), |_| Err(RedisError::WrongArity))
95    }
96}
97
98#[allow(clippy::not_unsafe_ptr_arg_deref)]
99pub fn decode_args(
100    ctx: *mut raw::RedisModuleCtx,
101    argv: *mut *mut raw::RedisModuleString,
102    argc: c_int,
103) -> Vec<RedisString> {
104    if argv.is_null() {
105        return Vec::new();
106    }
107    unsafe { slice::from_raw_parts(argv, argc as usize) }
108        .iter()
109        .map(|&arg| RedisString::new(NonNull::new(ctx), arg))
110        .collect()
111}
112
113///////////////////////////////////////////////////
114
115#[derive(Debug)]
116pub struct RedisString {
117    ctx: *mut raw::RedisModuleCtx,
118    pub inner: *mut raw::RedisModuleString,
119}
120
121impl RedisString {
122    pub(crate) fn take(mut self) -> *mut raw::RedisModuleString {
123        let inner = self.inner;
124        self.inner = std::ptr::null_mut();
125        inner
126    }
127
128    pub fn new(
129        ctx: Option<NonNull<raw::RedisModuleCtx>>,
130        inner: *mut raw::RedisModuleString,
131    ) -> Self {
132        let ctx = ctx.map_or(std::ptr::null_mut(), |v| v.as_ptr());
133        raw::string_retain_string(ctx, inner);
134        Self { ctx, inner }
135    }
136
137    /// In general, [RedisModuleString] is none atomic ref counted object.
138    /// So it is not safe to clone it if Redis GIL is not held.
139    /// [Self::safe_clone] gets a context reference which indicates that Redis GIL is held.
140    pub fn safe_clone(&self, _ctx: &Context) -> Self {
141        // RedisString are *not* atomic ref counted, so we must get a lock indicator to clone them.
142        // Alos notice that Redis allows us to create RedisModuleString with NULL context
143        // so we use [std::ptr::null_mut()] instead of the curren RedisString context.
144        // We do this because we can not promise the new RedisString will not outlive the current
145        // context and we want them to be independent.
146        raw::string_retain_string(ptr::null_mut(), self.inner);
147        Self {
148            ctx: ptr::null_mut(),
149            inner: self.inner,
150        }
151    }
152
153    #[allow(clippy::not_unsafe_ptr_arg_deref)]
154    pub fn create<T: Into<Vec<u8>>>(ctx: Option<NonNull<raw::RedisModuleCtx>>, s: T) -> Self {
155        let ctx = ctx.map_or(std::ptr::null_mut(), |v| v.as_ptr());
156        let str = CString::new(s).unwrap();
157        let inner = unsafe {
158            raw::RedisModule_CreateString.unwrap()(ctx, str.as_ptr(), str.as_bytes().len())
159        };
160
161        Self { ctx, inner }
162    }
163
164    /// Create a RedisString from a raw C string and length. The provided C String will be copied.
165    ///
166    /// # Safety
167    /// The caller must ensure that the provided pointer is valid and points to a memory region
168    /// that is at least `len` bytes long.
169    #[allow(clippy::not_unsafe_ptr_arg_deref)]
170    pub unsafe fn from_raw_parts(
171        ctx: Option<NonNull<raw::RedisModuleCtx>>,
172        s: *const c_char,
173        len: libc::size_t,
174    ) -> Self {
175        let ctx = ctx.map_or(std::ptr::null_mut(), |v| v.as_ptr());
176
177        let inner = unsafe { raw::RedisModule_CreateString.unwrap()(ctx, s, len) };
178
179        Self { ctx, inner }
180    }
181
182    #[allow(clippy::not_unsafe_ptr_arg_deref)]
183    pub fn create_from_slice(ctx: *mut raw::RedisModuleCtx, s: &[u8]) -> Self {
184        let inner = unsafe {
185            raw::RedisModule_CreateString.unwrap()(ctx, s.as_ptr().cast::<c_char>(), s.len())
186        };
187
188        Self { ctx, inner }
189    }
190
191    pub const fn from_redis_module_string(
192        ctx: *mut raw::RedisModuleCtx,
193        inner: *mut raw::RedisModuleString,
194    ) -> Self {
195        // Need to avoid string_retain_string
196        Self { ctx, inner }
197    }
198
199    #[allow(clippy::not_unsafe_ptr_arg_deref)]
200    pub fn from_ptr<'a>(ptr: *const raw::RedisModuleString) -> Result<&'a str, Utf8Error> {
201        str::from_utf8(Self::string_as_slice(ptr))
202    }
203
204    pub fn append(&mut self, s: &str) -> raw::Status {
205        raw::string_append_buffer(self.ctx, self.inner, s)
206    }
207
208    #[must_use]
209    pub fn len(&self) -> usize {
210        let mut len: usize = 0;
211        raw::string_ptr_len(self.inner, &mut len);
212        len
213    }
214
215    #[must_use]
216    pub fn is_empty(&self) -> bool {
217        let mut len: usize = 0;
218        raw::string_ptr_len(self.inner, &mut len);
219        len == 0
220    }
221
222    #[must_use]
223    pub fn as_cstr_ptr_and_len(&self) -> (*const c_char, usize) {
224        let mut len: usize = 0;
225        let ptr = raw::string_ptr_len(self.inner, &mut len);
226        (ptr, len)
227    }
228
229    pub fn try_as_str<'a>(&self) -> Result<&'a str, RedisError> {
230        Self::from_ptr(self.inner).map_err(|_| RedisError::Str("Couldn't parse as UTF-8 string"))
231    }
232
233    #[must_use]
234    pub fn as_slice(&self) -> &[u8] {
235        Self::string_as_slice(self.inner)
236    }
237
238    #[allow(clippy::not_unsafe_ptr_arg_deref)]
239    pub fn string_as_slice<'a>(ptr: *const raw::RedisModuleString) -> &'a [u8] {
240        let mut len: libc::size_t = 0;
241        let bytes = unsafe { raw::RedisModule_StringPtrLen.unwrap()(ptr, &mut len) };
242
243        unsafe { slice::from_raw_parts(bytes.cast::<u8>(), len) }
244    }
245
246    /// Performs lossy conversion of a `RedisString` into an owned `String. This conversion
247    /// will replace any invalid UTF-8 sequences with U+FFFD REPLACEMENT CHARACTER, which
248    /// looks like this: �.
249    ///
250    /// # Panics
251    ///
252    /// Will panic if `RedisModule_StringPtrLen` is missing in redismodule.h
253    #[must_use]
254    pub fn to_string_lossy(&self) -> String {
255        String::from_utf8_lossy(self.as_slice()).into_owned()
256    }
257
258    pub fn parse_unsigned_integer(&self) -> Result<u64, RedisError> {
259        let val = self.parse_integer()?;
260        u64::try_from(val)
261            .map_err(|_| RedisError::Str("Couldn't parse negative number as unsigned integer"))
262    }
263
264    pub fn parse_integer(&self) -> Result<i64, RedisError> {
265        let mut val: i64 = 0;
266        match raw::string_to_longlong(self.inner, &mut val) {
267            raw::Status::Ok => Ok(val),
268            raw::Status::Err => Err(RedisError::Str("Couldn't parse as integer")),
269        }
270    }
271
272    pub fn parse_float(&self) -> Result<f64, RedisError> {
273        let mut val: f64 = 0.0;
274        match raw::string_to_double(self.inner, &mut val) {
275            raw::Status::Ok => Ok(val),
276            raw::Status::Err => Err(RedisError::Str("Couldn't parse as float")),
277        }
278    }
279
280    // TODO: Redis allows storing and retrieving any arbitrary bytes.
281    // However rust's String and str can only store valid UTF-8.
282    // Implement these to allow non-utf8 bytes to be consumed:
283    // pub fn into_bytes(self) -> Vec<u8> {}
284    // pub fn as_bytes(&self) -> &[u8] {}
285}
286
287impl Drop for RedisString {
288    fn drop(&mut self) {
289        if !self.inner.is_null() {
290            unsafe {
291                raw::RedisModule_FreeString.unwrap()(self.ctx, self.inner);
292            }
293        }
294    }
295}
296
297impl PartialEq for RedisString {
298    fn eq(&self, other: &Self) -> bool {
299        self.cmp(other).is_eq()
300    }
301}
302
303impl Eq for RedisString {}
304
305impl PartialOrd for RedisString {
306    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
307        Some(self.cmp(other))
308    }
309}
310
311impl Ord for RedisString {
312    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
313        raw::string_compare(self.inner, other.inner)
314    }
315}
316
317impl core::hash::Hash for RedisString {
318    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
319        self.as_slice().hash(state);
320    }
321}
322
323impl Display for RedisString {
324    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
325        write!(f, "{}", self.to_string_lossy())
326    }
327}
328
329impl Borrow<str> for RedisString {
330    fn borrow(&self) -> &str {
331        // RedisString might not be UTF-8 safe
332        self.try_as_str().unwrap_or("<Invalid UTF-8 data>")
333    }
334}
335
336impl Clone for RedisString {
337    fn clone(&self) -> Self {
338        let inner =
339            // Redis allows us to create RedisModuleString with NULL context
340            // so we use [std::ptr::null_mut()] instead of the curren RedisString context.
341            // We do this because we can not promise the new RedisString will not outlive the current
342            // context and we want them to be independent.
343            unsafe { raw::RedisModule_CreateStringFromString.unwrap()(ptr::null_mut(), self.inner) };
344        Self::from_redis_module_string(ptr::null_mut(), inner)
345    }
346}
347
348impl From<RedisString> for String {
349    fn from(rs: RedisString) -> Self {
350        rs.to_string_lossy()
351    }
352}
353
354impl Deref for RedisString {
355    type Target = [u8];
356
357    fn deref(&self) -> &Self::Target {
358        self.as_slice()
359    }
360}
361
362impl From<RedisString> for Vec<u8> {
363    fn from(rs: RedisString) -> Self {
364        rs.as_slice().to_vec()
365    }
366}
367
368impl serde::Serialize for RedisString {
369    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
370    where
371        S: serde::Serializer,
372    {
373        serializer.serialize_bytes(self.as_slice())
374    }
375}
376
377struct RedisStringVisitor;
378
379impl<'de> serde::de::Visitor<'de> for RedisStringVisitor {
380    type Value = RedisString;
381
382    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
383        formatter.write_str("A bytes buffer")
384    }
385
386    fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
387    where
388        E: Error,
389    {
390        Ok(RedisString::create(None, v))
391    }
392
393    fn visit_seq<V>(self, mut visitor: V) -> Result<Self::Value, V::Error>
394    where
395        V: SeqAccess<'de>,
396    {
397        let mut v = if let Some(size_hint) = visitor.size_hint() {
398            Vec::with_capacity(size_hint)
399        } else {
400            Vec::new()
401        };
402        while let Some(elem) = visitor.next_element()? {
403            v.push(elem);
404        }
405
406        Ok(RedisString::create(None, v.as_slice()))
407    }
408}
409
410impl<'de> serde::Deserialize<'de> for RedisString {
411    fn deserialize<D>(deserializer: D) -> Result<RedisString, D::Error>
412    where
413        D: serde::Deserializer<'de>,
414    {
415        deserializer.deserialize_bytes(RedisStringVisitor)
416    }
417}
418
419///////////////////////////////////////////////////
420
421#[derive(Debug)]
422pub struct RedisBuffer {
423    buffer: *mut c_char,
424    len: usize,
425}
426
427impl RedisBuffer {
428    pub const fn new(buffer: *mut c_char, len: usize) -> Self {
429        Self { buffer, len }
430    }
431
432    pub fn to_string(&self) -> Result<String, FromUtf8Error> {
433        String::from_utf8(self.as_ref().to_vec())
434    }
435}
436
437impl AsRef<[u8]> for RedisBuffer {
438    fn as_ref(&self) -> &[u8] {
439        unsafe { slice::from_raw_parts(self.buffer as *const u8, self.len) }
440    }
441}
442
443impl Drop for RedisBuffer {
444    fn drop(&mut self) {
445        unsafe {
446            raw::RedisModule_Free.unwrap()(self.buffer.cast::<c_void>());
447        }
448    }
449}