Skip to main content

valkey_module/
redismodule.rs

1use std::borrow::Borrow;
2use std::ffi::CString;
3use std::fmt::Display;
4use std::ops::Deref;
5use std::os::raw::{c_char, c_int, c_void};
6use std::ptr::{null_mut, NonNull};
7use std::slice;
8use std::str;
9use std::str::Utf8Error;
10use std::string::FromUtf8Error;
11use std::{fmt, ptr};
12
13use serde::de::{Error, SeqAccess};
14
15pub use crate::raw;
16pub use crate::rediserror::ValkeyError;
17pub use crate::redisvalue::ValkeyValue;
18use crate::Context;
19
20/// A short-hand type that stores a [std::result::Result] with custom
21/// type and [RedisError].
22pub type ValkeyResult<T = ValkeyValue> = Result<T, ValkeyError>;
23/// A [RedisResult] with [ValkeyValue].
24pub type ValkeyValueResult = ValkeyResult<ValkeyValue>;
25
26impl From<ValkeyValue> for ValkeyValueResult {
27    fn from(v: ValkeyValue) -> Self {
28        Ok(v)
29    }
30}
31
32impl From<ValkeyError> for ValkeyValueResult {
33    fn from(v: ValkeyError) -> Self {
34        Err(v)
35    }
36}
37
38pub const VALKEY_OK: ValkeyValueResult = Ok(ValkeyValue::SimpleStringStatic("OK"));
39pub const TYPE_METHOD_VERSION: u64 = raw::REDISMODULE_TYPE_METHOD_VERSION as u64;
40pub const AUTH_HANDLED: i32 = raw::REDISMODULE_AUTH_HANDLED as i32;
41pub const AUTH_NOT_HANDLED: i32 = raw::REDISMODULE_AUTH_NOT_HANDLED as i32;
42
43pub trait NextArg {
44    fn next_arg(&mut self) -> Result<ValkeyString, ValkeyError>;
45    fn next_string(&mut self) -> Result<String, ValkeyError>;
46    fn next_str<'a>(&mut self) -> Result<&'a str, ValkeyError>;
47    fn next_i64(&mut self) -> Result<i64, ValkeyError>;
48    fn next_u64(&mut self) -> Result<u64, ValkeyError>;
49    fn next_f64(&mut self) -> Result<f64, ValkeyError>;
50    fn done(&mut self) -> Result<(), ValkeyError>;
51}
52
53impl<T> NextArg for T
54where
55    T: Iterator<Item = ValkeyString>,
56{
57    #[inline]
58    fn next_arg(&mut self) -> Result<ValkeyString, ValkeyError> {
59        self.next().ok_or(ValkeyError::WrongArity)
60    }
61
62    #[inline]
63    fn next_string(&mut self) -> Result<String, ValkeyError> {
64        self.next()
65            .map_or(Err(ValkeyError::WrongArity), |v| Ok(v.to_string_lossy()))
66    }
67
68    #[inline]
69    fn next_str<'a>(&mut self) -> Result<&'a str, ValkeyError> {
70        self.next()
71            .map_or(Err(ValkeyError::WrongArity), |v| v.try_as_str())
72    }
73
74    #[inline]
75    fn next_i64(&mut self) -> Result<i64, ValkeyError> {
76        self.next()
77            .map_or(Err(ValkeyError::WrongArity), |v| v.parse_integer())
78    }
79
80    #[inline]
81    fn next_u64(&mut self) -> Result<u64, ValkeyError> {
82        self.next()
83            .map_or(Err(ValkeyError::WrongArity), |v| v.parse_unsigned_integer())
84    }
85
86    #[inline]
87    fn next_f64(&mut self) -> Result<f64, ValkeyError> {
88        self.next()
89            .map_or(Err(ValkeyError::WrongArity), |v| v.parse_float())
90    }
91
92    /// Return an error if there are any more arguments
93    #[inline]
94    fn done(&mut self) -> Result<(), ValkeyError> {
95        self.next().map_or(Ok(()), |_| Err(ValkeyError::WrongArity))
96    }
97}
98
99#[allow(clippy::not_unsafe_ptr_arg_deref)]
100pub fn decode_args(
101    ctx: *mut raw::RedisModuleCtx,
102    argv: *mut *mut raw::RedisModuleString,
103    argc: c_int,
104) -> Vec<ValkeyString> {
105    if argv.is_null() {
106        return Vec::new();
107    }
108    unsafe { slice::from_raw_parts(argv, argc as usize) }
109        .iter()
110        .map(|&arg| ValkeyString::new(NonNull::new(ctx), arg))
111        .collect()
112}
113
114///////////////////////////////////////////////////
115
116#[derive(Debug)]
117pub struct ValkeyString {
118    ctx: *mut raw::RedisModuleCtx,
119    pub inner: *mut raw::RedisModuleString,
120}
121
122impl ValkeyString {
123    pub(crate) fn take(mut self) -> *mut raw::RedisModuleString {
124        let inner = self.inner;
125        self.inner = std::ptr::null_mut();
126        inner
127    }
128
129    pub fn new(
130        ctx: Option<NonNull<raw::RedisModuleCtx>>,
131        inner: *mut raw::RedisModuleString,
132    ) -> Self {
133        let ctx = ctx.map_or(std::ptr::null_mut(), |v| v.as_ptr());
134        raw::string_retain_string(ctx, inner);
135        Self { ctx, inner }
136    }
137
138    /// In general, [RedisModuleString] is none atomic ref counted object.
139    /// So it is not safe to clone it if Valkey GIL is not held.
140    /// [Self::safe_clone] gets a context reference which indicates that Valkey GIL is held.
141    pub fn safe_clone(&self, _ctx: &Context) -> Self {
142        // RedisString are *not* atomic ref counted, so we must get a lock indicator to clone them.
143        // Alos notice that Valkey allows us to create RedisModuleString with NULL context
144        // so we use [std::ptr::null_mut()] instead of the curren RedisString context.
145        // We do this because we can not promise the new RedisString will not outlive the current
146        // context and we want them to be independent.
147        raw::string_retain_string(ptr::null_mut(), self.inner);
148        Self {
149            ctx: ptr::null_mut(),
150            inner: self.inner,
151        }
152    }
153
154    #[allow(clippy::not_unsafe_ptr_arg_deref)]
155    pub fn create<T: Into<Vec<u8>>>(ctx: Option<NonNull<raw::RedisModuleCtx>>, s: T) -> Self {
156        let ctx = ctx.map_or(std::ptr::null_mut(), |v| v.as_ptr());
157        let str = CString::new(s).unwrap();
158        let inner = unsafe {
159            raw::RedisModule_CreateString.unwrap()(ctx, str.as_ptr(), str.as_bytes().len())
160        };
161
162        Self { ctx, inner }
163    }
164
165    #[allow(clippy::not_unsafe_ptr_arg_deref)]
166    pub fn create_from_slice(ctx: *mut raw::RedisModuleCtx, s: &[u8]) -> Self {
167        let inner = unsafe {
168            raw::RedisModule_CreateString.unwrap()(ctx, s.as_ptr().cast::<c_char>(), s.len())
169        };
170
171        Self { ctx, inner }
172    }
173
174    /// Creates a ValkeyString from a &str and retains it.  This is useful in cases where Modules need to pass ownership of a ValkeyString to the core engine without it being freed when we drop a ValkeyString
175    pub fn create_and_retain(arg: &str) -> ValkeyString {
176        let arg = ValkeyString::create(None, arg);
177        raw::string_retain_string(null_mut(), arg.inner);
178        arg
179    }
180
181    pub const fn from_redis_module_string(
182        ctx: *mut raw::RedisModuleCtx,
183        inner: *mut raw::RedisModuleString,
184    ) -> Self {
185        // Need to avoid string_retain_string
186        Self { ctx, inner }
187    }
188
189    #[allow(clippy::not_unsafe_ptr_arg_deref)]
190    pub fn from_ptr<'a>(ptr: *const raw::RedisModuleString) -> Result<&'a str, Utf8Error> {
191        str::from_utf8(Self::string_as_slice(ptr))
192    }
193
194    pub fn append(&mut self, s: &str) -> raw::Status {
195        raw::string_append_buffer(self.ctx, self.inner, s)
196    }
197
198    #[must_use]
199    pub fn len(&self) -> usize {
200        let mut len: usize = 0;
201        raw::string_ptr_len(self.inner, &mut len);
202        len
203    }
204
205    #[must_use]
206    pub fn is_empty(&self) -> bool {
207        let mut len: usize = 0;
208        raw::string_ptr_len(self.inner, &mut len);
209        len == 0
210    }
211
212    pub fn try_as_str<'a>(&self) -> Result<&'a str, ValkeyError> {
213        Self::from_ptr(self.inner).map_err(|_| ValkeyError::Str("Couldn't parse as UTF-8 string"))
214    }
215
216    #[must_use]
217    pub fn as_slice(&self) -> &[u8] {
218        Self::string_as_slice(self.inner)
219    }
220
221    #[allow(clippy::not_unsafe_ptr_arg_deref)]
222    pub fn string_as_slice<'a>(ptr: *const raw::RedisModuleString) -> &'a [u8] {
223        let mut len: libc::size_t = 0;
224        let bytes = unsafe { raw::RedisModule_StringPtrLen.unwrap()(ptr, &mut len) };
225
226        unsafe { slice::from_raw_parts(bytes.cast::<u8>(), len) }
227    }
228
229    /// Performs lossy conversion of a `RedisString` into an owned `String. This conversion
230    /// will replace any invalid UTF-8 sequences with U+FFFD REPLACEMENT CHARACTER, which
231    /// looks like this: �.
232    ///
233    /// # Panics
234    ///
235    /// Will panic if `RedisModule_StringPtrLen` is missing in redismodule.h
236    #[must_use]
237    pub fn to_string_lossy(&self) -> String {
238        String::from_utf8_lossy(self.as_slice()).into_owned()
239    }
240
241    pub fn parse_unsigned_integer(&self) -> Result<u64, ValkeyError> {
242        let mut val: u64 = 0;
243        match raw::string_to_ulonglong(self.inner, &mut val) {
244            raw::Status::Ok => Ok(val),
245            raw::Status::Err => Err(ValkeyError::Str("Couldn't parse as unsigned integer")),
246        }
247    }
248
249    pub fn parse_integer(&self) -> Result<i64, ValkeyError> {
250        let mut val: i64 = 0;
251        match raw::string_to_longlong(self.inner, &mut val) {
252            raw::Status::Ok => Ok(val),
253            raw::Status::Err => Err(ValkeyError::Str("Couldn't parse as integer")),
254        }
255    }
256
257    pub fn parse_float(&self) -> Result<f64, ValkeyError> {
258        let mut val: f64 = 0.0;
259        match raw::string_to_double(self.inner, &mut val) {
260            raw::Status::Ok => Ok(val),
261            raw::Status::Err => Err(ValkeyError::Str("Couldn't parse as float")),
262        }
263    }
264
265    // TODO: Valkey allows storing and retrieving any arbitrary bytes.
266    // However rust's String and str can only store valid UTF-8.
267    // Implement these to allow non-utf8 bytes to be consumed:
268    // pub fn into_bytes(self) -> Vec<u8> {}
269    // pub fn as_bytes(&self) -> &[u8] {}
270}
271
272impl Drop for ValkeyString {
273    fn drop(&mut self) {
274        if !self.inner.is_null() {
275            unsafe {
276                raw::RedisModule_FreeString.unwrap()(self.ctx, self.inner);
277            }
278        }
279    }
280}
281
282impl PartialEq for ValkeyString {
283    fn eq(&self, other: &Self) -> bool {
284        self.cmp(other).is_eq()
285    }
286}
287
288impl Eq for ValkeyString {}
289
290impl PartialOrd for ValkeyString {
291    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
292        Some(self.cmp(other))
293    }
294}
295
296impl Ord for ValkeyString {
297    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
298        raw::string_compare(self.inner, other.inner)
299    }
300}
301
302impl core::hash::Hash for ValkeyString {
303    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
304        self.as_slice().hash(state);
305    }
306}
307
308impl Display for ValkeyString {
309    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
310        write!(f, "{}", self.to_string_lossy())
311    }
312}
313
314impl Borrow<str> for ValkeyString {
315    fn borrow(&self) -> &str {
316        // RedisString might not be UTF-8 safe
317        self.try_as_str().unwrap_or("<Invalid UTF-8 data>")
318    }
319}
320
321impl Clone for ValkeyString {
322    fn clone(&self) -> Self {
323        let inner =
324            // Valkey allows us to create RedisModuleString with NULL context
325            // so we use [std::ptr::null_mut()] instead of the curren RedisString context.
326            // We do this because we can not promise the new RedisString will not outlive the current
327            // context and we want them to be independent.
328            unsafe { raw::RedisModule_CreateStringFromString.unwrap()(ptr::null_mut(), self.inner) };
329        Self::from_redis_module_string(ptr::null_mut(), inner)
330    }
331}
332
333impl From<ValkeyString> for String {
334    fn from(rs: ValkeyString) -> Self {
335        rs.to_string_lossy()
336    }
337}
338
339impl Deref for ValkeyString {
340    type Target = [u8];
341
342    fn deref(&self) -> &Self::Target {
343        self.as_slice()
344    }
345}
346
347impl From<ValkeyString> for Vec<u8> {
348    fn from(rs: ValkeyString) -> Self {
349        rs.as_slice().to_vec()
350    }
351}
352
353impl serde::Serialize for ValkeyString {
354    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
355    where
356        S: serde::Serializer,
357    {
358        serializer.serialize_bytes(self.as_slice())
359    }
360}
361
362struct RedisStringVisitor;
363
364impl<'de> serde::de::Visitor<'de> for RedisStringVisitor {
365    type Value = ValkeyString;
366
367    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
368        formatter.write_str("A bytes buffer")
369    }
370
371    fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
372    where
373        E: Error,
374    {
375        Ok(ValkeyString::create(None, v))
376    }
377
378    fn visit_seq<V>(self, mut visitor: V) -> Result<Self::Value, V::Error>
379    where
380        V: SeqAccess<'de>,
381    {
382        let mut v = if let Some(size_hint) = visitor.size_hint() {
383            Vec::with_capacity(size_hint)
384        } else {
385            Vec::new()
386        };
387        while let Some(elem) = visitor.next_element()? {
388            v.push(elem);
389        }
390
391        Ok(ValkeyString::create(None, v.as_slice()))
392    }
393}
394
395impl<'de> serde::Deserialize<'de> for ValkeyString {
396    fn deserialize<D>(deserializer: D) -> Result<ValkeyString, D::Error>
397    where
398        D: serde::Deserializer<'de>,
399    {
400        deserializer.deserialize_bytes(RedisStringVisitor)
401    }
402}
403
404///////////////////////////////////////////////////
405
406#[derive(Debug)]
407pub struct RedisBuffer {
408    buffer: *mut c_char,
409    len: usize,
410}
411
412impl RedisBuffer {
413    pub const fn new(buffer: *mut c_char, len: usize) -> Self {
414        Self { buffer, len }
415    }
416
417    pub fn to_string(&self) -> Result<String, FromUtf8Error> {
418        String::from_utf8(self.as_ref().to_vec())
419    }
420}
421
422impl AsRef<[u8]> for RedisBuffer {
423    fn as_ref(&self) -> &[u8] {
424        unsafe { slice::from_raw_parts(self.buffer as *const u8, self.len) }
425    }
426}
427
428impl Drop for RedisBuffer {
429    fn drop(&mut self) {
430        unsafe {
431            raw::RedisModule_Free.unwrap()(self.buffer.cast::<c_void>());
432        }
433    }
434}
435
436#[cfg(test)]
437mod tests {
438    use super::*;
439
440    #[test]
441    fn safe_clone_keeps_string_alive_after_original_is_dropped() {
442        let context = Context::test();
443        let original = ValkeyString::test("value");
444        let cloned = original.safe_clone(&context);
445
446        assert_eq!(cloned.inner, original.inner);
447        drop(original);
448
449        assert_eq!(cloned.as_slice(), b"value");
450    }
451
452    #[test]
453    fn create_and_retain_keeps_the_transferred_reference_alive() {
454        let _context = Context::test();
455        let string = ValkeyString::create_and_retain("value");
456        let inner = string.inner;
457
458        assert_eq!(string.as_slice(), b"value");
459        drop(string);
460
461        // SAFETY: `create_and_retain` creates a second shim-backed reference for the caller to
462        // transfer to Valkey. This releases that reference after the test's owner was dropped.
463        unsafe {
464            raw::RedisModule_FreeString.unwrap()(ptr::null_mut(), inner);
465        }
466    }
467
468    #[test]
469    fn create_from_slice_preserves_binary_data() {
470        let _context = Context::test();
471        let string = ValkeyString::create_from_slice(ptr::null_mut(), &[0, 0xff, b'a']);
472
473        assert_eq!(string.as_slice(), &[0, 0xff, b'a']);
474    }
475
476    #[test]
477    fn empty_string_is_empty() {
478        let string = ValkeyString::test("");
479
480        assert!(string.is_empty());
481    }
482
483    #[test]
484    fn invalid_utf8_reports_an_error_and_converts_lossily() {
485        let string = ValkeyString::test([b'a', 0xff]);
486
487        assert!(matches!(
488            string.try_as_str(),
489            Err(ValkeyError::Str("Couldn't parse as UTF-8 string"))
490        ));
491        assert_eq!(string.to_string_lossy(), "a�");
492    }
493
494    #[test]
495    fn parses_signed_integer_boundaries() {
496        let minimum = ValkeyString::test(i64::MIN.to_string());
497        let maximum = ValkeyString::test(i64::MAX.to_string());
498
499        assert_eq!(
500            minimum
501                .parse_integer()
502                .expect("minimum signed integer should parse"),
503            i64::MIN
504        );
505        assert_eq!(
506            maximum
507                .parse_integer()
508                .expect("maximum signed integer should parse"),
509            i64::MAX
510        );
511    }
512
513    #[test]
514    fn rejects_invalid_signed_and_unsigned_integers() {
515        let signed_overflow = ValkeyString::test("9223372036854775808");
516        let negative_unsigned = ValkeyString::test("-1");
517        let unsigned_overflow = ValkeyString::test("18446744073709551616");
518
519        assert!(matches!(
520            signed_overflow.parse_integer(),
521            Err(ValkeyError::Str("Couldn't parse as integer"))
522        ));
523        assert!(matches!(
524            negative_unsigned.parse_unsigned_integer(),
525            Err(ValkeyError::Str("Couldn't parse as unsigned integer"))
526        ));
527        assert!(matches!(
528            unsigned_overflow.parse_unsigned_integer(),
529            Err(ValkeyError::Str("Couldn't parse as unsigned integer"))
530        ));
531    }
532
533    #[test]
534    fn rejects_invalid_float() {
535        let string = ValkeyString::test("not-a-float");
536
537        assert!(matches!(
538            string.parse_float(),
539            Err(ValkeyError::Str("Couldn't parse as float"))
540        ));
541    }
542
543    #[test]
544    fn next_arg_helpers_consume_values_of_each_supported_type() {
545        let mut args = vec![
546            ValkeyString::test([0, 0xff]),
547            ValkeyString::test("string"),
548            ValkeyString::test("text"),
549            ValkeyString::test("-42"),
550            ValkeyString::test(u64::MAX.to_string()),
551            ValkeyString::test("42.5"),
552        ]
553        .into_iter();
554
555        assert_eq!(
556            args.next_arg()
557                .expect("binary argument should be returned")
558                .as_slice(),
559            &[0, 0xff]
560        );
561        assert_eq!(
562            args.next_string()
563                .expect("string argument should be returned"),
564            "string"
565        );
566        assert!(args.next_str().is_ok());
567        assert_eq!(args.next_i64().expect("integer should parse"), -42);
568        assert_eq!(
569            args.next_u64().expect("unsigned integer should parse"),
570            u64::MAX
571        );
572        assert_eq!(args.next_f64().expect("float should parse"), 42.5);
573        assert!(args.done().is_ok());
574    }
575
576    #[test]
577    fn next_arg_helpers_report_wrong_arity_for_empty_iterator() {
578        let mut args = Vec::<ValkeyString>::new().into_iter();
579
580        assert!(matches!(args.next_arg(), Err(ValkeyError::WrongArity)));
581        assert!(matches!(args.next_string(), Err(ValkeyError::WrongArity)));
582        assert!(matches!(args.next_str(), Err(ValkeyError::WrongArity)));
583        assert!(matches!(args.next_i64(), Err(ValkeyError::WrongArity)));
584        assert!(matches!(args.next_u64(), Err(ValkeyError::WrongArity)));
585        assert!(matches!(args.next_f64(), Err(ValkeyError::WrongArity)));
586    }
587
588    #[test]
589    fn next_str_rejects_invalid_utf8_while_next_string_is_lossy() {
590        let mut string_args = vec![ValkeyString::test([b'a', 0xff])].into_iter();
591        let mut str_args = vec![ValkeyString::test([b'a', 0xff])].into_iter();
592
593        assert_eq!(
594            string_args
595                .next_string()
596                .expect("next_string should use a lossy conversion"),
597            "a�"
598        );
599        assert!(matches!(
600            str_args.next_str(),
601            Err(ValkeyError::Str("Couldn't parse as UTF-8 string"))
602        ));
603    }
604
605    #[test]
606    fn numeric_next_arg_helpers_report_parse_errors() {
607        let mut integer_args = vec![ValkeyString::test("not-an-integer")].into_iter();
608        let mut unsigned_args = vec![ValkeyString::test("-1")].into_iter();
609        let mut float_args = vec![ValkeyString::test("not-a-float")].into_iter();
610
611        assert!(matches!(
612            integer_args.next_i64(),
613            Err(ValkeyError::Str("Couldn't parse as integer"))
614        ));
615        assert!(matches!(
616            unsigned_args.next_u64(),
617            Err(ValkeyError::Str("Couldn't parse as unsigned integer"))
618        ));
619        assert!(matches!(
620            float_args.next_f64(),
621            Err(ValkeyError::Str("Couldn't parse as float"))
622        ));
623    }
624
625    #[test]
626    fn done_rejects_remaining_arguments() {
627        let mut args = vec![ValkeyString::test("extra")].into_iter();
628
629        assert!(matches!(args.done(), Err(ValkeyError::WrongArity)));
630    }
631}