Handle Short Reads (#178)

* Handle Short Reads

* Remove 6.0.15 from the versions supporting Short Read for modules

* Move away check_minimal_version_for_short_read
+ Improve get_redis_version

* Use clear error message in get_redis_version

* Trivial commit to trigger CI

* Trivial commit to hopefully trigger CI

* Update Cargo.toml

* Move get_redis_version stuff out (done in PR #184)

* Fix memory leak in redis_module::raw::load_string (asan)

* WIP - code review by Gavrie - remove dead code

* Trivial commit to trigger CI

* Revert "WIP - code review by Gavrie - remove dead code"
This is not dead code (used by RedisJSON)
This reverts commit 879815f.

* load with generics

* Add Context api set_module_options

* Remove dead code (used by aux_load removed from RedisJSON)

This reverts commit c32d985.

Co-authored-by: Guy Korland <[email protected]>
Co-authored-by: Rafi Einstein <[email protected]>

Author: 3 people

src/context/mod.rs

@@ -4,6 +4,7 @@ use std::ptr;
 
 use crate::key::{RedisKey, RedisKeyWritable};
 use crate::raw;
+use crate::raw::ModuleOptions;
 use crate::LogLevel;
 use crate::{RedisError, RedisResult, RedisString, RedisValue};
 
@@ -263,4 +264,8 @@ impl Context {
     ) -> raw::Status {
         raw::notify_keyspace_event(self.ctx, event_type, event, keyname)
     }
+
+    pub fn set_module_options(&self, options: ModuleOptions) {
+        unsafe { raw::RedisModule_SetModuleOptions.unwrap()(self.ctx, options.bits()) };
+    }
 }

src/error.rs

@@ -1,3 +1,4 @@
+use crate::RedisError;
 use std::error;
 use std::fmt;
 use std::fmt::Display;
@@ -15,6 +16,12 @@ impl Error {
     }
 }
 
+impl From<RedisError> for Error {
+    fn from(err: RedisError) -> Error {
+        Error::generic(err.to_string().as_str())
+    }
+}
+
 impl From<std::string::FromUtf8Error> for Error {
     fn from(err: std::string::FromUtf8Error) -> Error {
         Error::FromUtf8(err)

src/raw.rs

@@ -6,15 +6,17 @@ extern crate enum_primitive_derive;
 extern crate libc;
 extern crate num_traits;
 
-use bitflags::bitflags;
-use enum_primitive_derive::Primitive;
-use libc::size_t;
-use num_traits::FromPrimitive;
 use std::ffi::CString;
 use std::os::raw::{c_char, c_double, c_int, c_long, c_longlong};
 use std::ptr;
 use std::slice;
 
+use bitflags::bitflags;
+use enum_primitive_derive::Primitive;
+use libc::size_t;
+use num_traits::FromPrimitive;
+
+use crate::error::Error;
 pub use crate::redisraw::bindings::*;
 use crate::RedisString;
 use crate::{RedisBuffer, RedisError};
@@ -26,6 +28,13 @@ bitflags! {
     }
 }
 
+bitflags! {
+    pub struct ModuleOptions: c_int {
+        const HANDLE_IO_ERRORS = REDISMODULE_OPTIONS_HANDLE_IO_ERRORS as c_int;
+        const NO_IMPLICIT_SIGNAL_MODIFIED = REDISMODULE_OPTION_NO_IMPLICIT_SIGNAL_MODIFIED as c_int;
+    }
+}
+
 #[derive(Primitive, Debug, PartialEq)]
 pub enum KeyType {
     Empty = REDISMODULE_KEYTYPE_EMPTY,
@@ -445,30 +454,43 @@ pub fn replicate_verbatim(ctx: *mut RedisModuleCtx) -> Status {
     unsafe { RedisModule_ReplicateVerbatim.unwrap()(ctx).into() }
 }
 
+fn load<F, T>(rdb: *mut RedisModuleIO, f: F) -> Result<T, Error>
+where
+    F: FnOnce(*mut RedisModuleIO) -> T,
+{
+    let res = f(rdb);
+    if is_io_error(rdb) {
+        Err(RedisError::short_read().into())
+    } else {
+        Ok(res)
+    }
+}
+
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
-pub fn load_unsigned(rdb: *mut RedisModuleIO) -> u64 {
-    unsafe { RedisModule_LoadUnsigned.unwrap()(rdb) }
+pub fn load_unsigned(rdb: *mut RedisModuleIO) -> Result<u64, Error> {
+    unsafe { load(rdb, |rdb| RedisModule_LoadUnsigned.unwrap()(rdb)) }
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
-pub fn load_signed(rdb: *mut RedisModuleIO) -> i64 {
-    unsafe { RedisModule_LoadSigned.unwrap()(rdb) }
+pub fn load_signed(rdb: *mut RedisModuleIO) -> Result<i64, Error> {
+    unsafe { load(rdb, |rdb| RedisModule_LoadSigned.unwrap()(rdb)) }
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
-pub fn load_string(rdb: *mut RedisModuleIO) -> String {
-    let p = unsafe { RedisModule_LoadString.unwrap()(rdb) };
-    RedisString::from_ptr(p)
-        .expect("UTF8 encoding error in load string")
-        .to_string()
+pub fn load_string(rdb: *mut RedisModuleIO) -> Result<RedisString, Error> {
+    let p = unsafe { load(rdb, |rdb| RedisModule_LoadString.unwrap()(rdb))? };
+    let ctx = unsafe { RedisModule_GetContextFromIO.unwrap()(rdb) };
+    Ok(RedisString::from_redis_module_string(ctx, p))
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
-pub fn load_string_buffer(rdb: *mut RedisModuleIO) -> RedisBuffer {
+pub fn load_string_buffer(rdb: *mut RedisModuleIO) -> Result<RedisBuffer, Error> {
     unsafe {
         let mut len = 0;
-        let buffer = RedisModule_LoadStringBuffer.unwrap()(rdb, &mut len);
-        RedisBuffer::new(buffer, len)
+        let buffer = load(rdb, |rdb| {
+            RedisModule_LoadStringBuffer.unwrap()(rdb, &mut len)
+        })?;
+        Ok(RedisBuffer::new(buffer, len))
     }
 }
 
@@ -494,13 +516,13 @@ pub fn replicate(ctx: *mut RedisModuleCtx, command: &str, args: &[&str]) -> Stat
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
-pub fn load_double(rdb: *mut RedisModuleIO) -> f64 {
-    unsafe { RedisModule_LoadDouble.unwrap()(rdb) }
+pub fn load_double(rdb: *mut RedisModuleIO) -> Result<f64, Error> {
+    unsafe { load(rdb, |rdb| RedisModule_LoadDouble.unwrap()(rdb)) }
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
-pub fn load_float(rdb: *mut RedisModuleIO) -> f32 {
-    unsafe { RedisModule_LoadFloat.unwrap()(rdb) }
+pub fn load_float(rdb: *mut RedisModuleIO) -> Result<f32, Error> {
+    unsafe { load(rdb, |rdb| RedisModule_LoadFloat.unwrap()(rdb)) }
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
@@ -584,3 +606,8 @@ pub fn get_keyspace_events() -> NotifyEvent {
         NotifyEvent::from_bits_truncate(events)
     }
 }
+
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+pub fn is_io_error(rdb: *mut RedisModuleIO) -> bool {
+    unsafe { RedisModule_IsIOError.unwrap()(rdb) != 0 }
+}

src/rediserror.rs

@@ -14,6 +14,10 @@ impl RedisError {
     pub fn nonexistent_key() -> Self {
         RedisError::Str("ERR could not perform this operation on a key that doesn't exist")
     }
+
+    pub fn short_read() -> Self {
+        RedisError::Str("ERR short read or OOM loading DB")
+    }
 }
 
 impl<T: std::error::Error> From<T> for RedisError {

src/redismodule.rs

@@ -50,7 +50,7 @@ where
     }
 
     #[inline]
-   fn next_i64(&mut self) -> Result<i64, RedisError> {
+    fn next_i64(&mut self) -> Result<i64, RedisError> {
         self.next()
             .map_or(Err(RedisError::WrongArity), |v| v.parse_integer())
     }
@@ -109,6 +109,14 @@ impl RedisString {
         RedisString { ctx, inner }
     }
 
+    pub fn from_redis_module_string(
+        ctx: *mut raw::RedisModuleCtx,
+        inner: *mut raw::RedisModuleString,
+    ) -> RedisString {
+        // Need to avoid string_retain_string
+        RedisString { ctx, inner }
+    }
+
     #[allow(clippy::not_unsafe_ptr_arg_deref)]
     pub fn from_ptr<'a>(ptr: *const raw::RedisModuleString) -> Result<&'a str, Utf8Error> {
         let mut len: libc::size_t = 0;