Sanitize dump payload: fail RESTORE if memory allocation fails

When RDB input attempts to make a huge memory allocation that fails,
RESTORE should fail gracefully rather than die with panic

9 files changed, 199 insertions, 83 deletions

diff --git a/src/dict.c b/src/dict.c
index cca5aa921..4736dacd5 100644
--- a/src/dict.c
+++ b/src/dict.c
@@ -143,9 +143,13 @@ int dictResize(dict *d)
     return dictExpand(d, minimal);
 }
 
-/* Expand or create the hash table */
-int dictExpand(dict *d, unsigned long size)
+/* Expand or create the hash table,
+ * when malloc_failed is non-NULL, it'll avoid panic if malloc fails (in which case it'll be set to 1).
+ * Returns DICT_OK if expand was performed, and DICT_ERR if skipped. */
+int _dictExpand(dict *d, unsigned long size, int* malloc_failed)
 {
+    if (malloc_failed) *malloc_failed = 0;
+
     /* the size is invalid if it is smaller than the number of
      * elements already inside the hash table */
     if (dictIsRehashing(d) || d->ht[0].used > size)
@@ -160,7 +164,14 @@ int dictExpand(dict *d, unsigned long size)
     /* Allocate the new hash table and initialize all pointers to NULL */
     n.size = realsize;
     n.sizemask = realsize-1;
-    n.table = zcalloc(realsize*sizeof(dictEntry*));
+    if (malloc_failed) {
+        n.table = ztrycalloc(realsize*sizeof(dictEntry*));
+        *malloc_failed = n.table == NULL;
+        if (*malloc_failed)
+            return DICT_ERR;
+    } else
+        n.table = zcalloc(realsize*sizeof(dictEntry*));
+
     n.used = 0;
 
     /* Is this the first initialization? If so it's not really a rehashing
@@ -176,6 +187,18 @@ int dictExpand(dict *d, unsigned long size)
     return DICT_OK;
 }
 
+/* return DICT_ERR if expand was not performed */
+int dictExpand(dict *d, unsigned long size) {
+    return _dictExpand(d, size, NULL);
+}
+
+/* return DICT_ERR if expand failed due to memory allocation failure */
+int dictTryExpand(dict *d, unsigned long size) {
+    int malloc_failed;
+    _dictExpand(d, size, &malloc_failed);
+    return malloc_failed? DICT_ERR : DICT_OK;
+}
+
 /* Performs N steps of incremental rehashing. Returns 1 if there are still
  * keys to move from the old to the new hash table, otherwise 0 is returned.
  *
diff --git a/src/dict.h b/src/dict.h
index fca924abe..f7515e905 100644
--- a/src/dict.h
+++ b/src/dict.h
@@ -151,6 +151,7 @@ typedef void (dictScanBucketFunction)(void *privdata, dictEntry **bucketref);
 /* API */
 dict *dictCreate(dictType *type, void *privDataPtr);
 int dictExpand(dict *d, unsigned long size);
+int dictTryExpand(dict *d, unsigned long size);
 int dictAdd(dict *d, void *key, void *val);
 dictEntry *dictAddRaw(dict *d, void *key, dictEntry **existing);
 dictEntry *dictAddOrFind(dict *d, void *key);
diff --git a/src/rdb.c b/src/rdb.c
index e8501ee50..0878be831 100644
--- a/src/rdb.c
+++ b/src/rdb.c
@@ -387,14 +387,22 @@ void *rdbLoadLzfStringObject(rio *rdb, int flags, size_t *lenptr) {
 
     if ((clen = rdbLoadLen(rdb,NULL)) == RDB_LENERR) return NULL;
     if ((len = rdbLoadLen(rdb,NULL)) == RDB_LENERR) return NULL;
-    if ((c = zmalloc(clen)) == NULL) goto err;
+    if ((c = ztrymalloc(clen)) == NULL) {
+        serverLog(server.loading? LL_WARNING: LL_VERBOSE, "rdbLoadLzfStringObject failed allocating %llu bytes", (unsigned long long)clen);
+        goto err;
+    }
 
     /* Allocate our target according to the uncompressed size. */
     if (plain) {
-        val = zmalloc(len);
+        val = ztrymalloc(len);
     } else {
-        val = sdsnewlen(SDS_NOINIT,len);
+        val = sdstrynewlen(SDS_NOINIT,len);
+    }
+    if (!val) {
+        serverLog(server.loading? LL_WARNING: LL_VERBOSE, "rdbLoadLzfStringObject failed allocating %llu bytes", (unsigned long long)len);
+        goto err;
     }
+
     if (lenptr) *lenptr = len;
 
     /* Load the compressed representation and uncompress it to target. */
@@ -522,7 +530,11 @@ void *rdbGenericLoadStringObject(rio *rdb, int flags, size_t *lenptr) {
     }
 
     if (plain || sds) {
-        void *buf = plain ? zmalloc(len) : sdsnewlen(SDS_NOINIT,len);
+        void *buf = plain ? ztrymalloc(len) : sdstrynewlen(SDS_NOINIT,len);
+        if (!buf) {
+            serverLog(server.loading? LL_WARNING: LL_VERBOSE, "rdbGenericLoadStringObject failed allocating %llu bytes", len);
+            return NULL;
+        }
         if (lenptr) *lenptr = len;
         if (len && rioRead(rdb,buf,len) == 0) {
             if (plain)
@@ -1545,8 +1557,11 @@ robj *rdbLoadObject(int rdbtype, rio *rdb, sds key) {
             o = createSetObject();
             /* It's faster to expand the dict to the right size asap in order
              * to avoid rehashing */
-            if (len > DICT_HT_INITIAL_SIZE)
-                dictExpand(o->ptr,len);
+            if (len > DICT_HT_INITIAL_SIZE && dictTryExpand(o->ptr,len) != DICT_OK) {
+                rdbReportCorruptRDB("OOM in dictTryExpand %llu", (unsigned long long)len);
+                decrRefCount(o);
+                return NULL;
+            }
         } else {
             o = createIntsetObject();
         }
@@ -1574,7 +1589,12 @@ robj *rdbLoadObject(int rdbtype, rio *rdb, sds key) {
                     }
                 } else {
                     setTypeConvert(o,OBJ_ENCODING_HT);
-                    dictExpand(o->ptr,len);
+                    if (dictTryExpand(o->ptr,len) != DICT_OK) {
+                        rdbReportCorruptRDB("OOM in dictTryExpand %llu", (unsigned long long)len);
+                        sdsfree(sdsele);
+                        decrRefCount(o);
+                        return NULL;
+                    }
                 }
             }
 
@@ -1601,8 +1621,11 @@ robj *rdbLoadObject(int rdbtype, rio *rdb, sds key) {
         o = createZsetObject();
         zs = o->ptr;
 
-        if (zsetlen > DICT_HT_INITIAL_SIZE)
-            dictExpand(zs->dict,zsetlen);
+        if (zsetlen > DICT_HT_INITIAL_SIZE && dictTryExpand(zs->dict,zsetlen) != DICT_OK) {
+            rdbReportCorruptRDB("OOM in dictTryExpand %llu", (unsigned long long)zsetlen);
+            decrRefCount(o);
+            return NULL;
+        }
 
         /* Load every single element of the sorted set. */
         while(zsetlen--) {
@@ -1723,8 +1746,13 @@ robj *rdbLoadObject(int rdbtype, rio *rdb, sds key) {
             dupSearchDict = NULL;
         }
 
-        if (o->encoding == OBJ_ENCODING_HT && len > DICT_HT_INITIAL_SIZE)
-            dictExpand(o->ptr,len);
+        if (o->encoding == OBJ_ENCODING_HT && len > DICT_HT_INITIAL_SIZE) {
+            if (dictTryExpand(o->ptr,len) != DICT_OK) {
+                rdbReportCorruptRDB("OOM in dictTryExpand %llu", (unsigned long long)len);
+                decrRefCount(o);
+                return NULL;
+            }
+        }
 
         /* Load remaining fields and values into the hash table */
         while (o->encoding == OBJ_ENCODING_HT && len > 0) {
@@ -1823,9 +1851,9 @@ robj *rdbLoadObject(int rdbtype, rio *rdb, sds key) {
                         zl = ziplistPush(zl, vstr, vlen, ZIPLIST_TAIL);
 
                         /* search for duplicate records */
-                        sds field = sdsnewlen(fstr, flen);
-                        if (dictAdd(dupSearchDict, field, NULL) != DICT_OK) {
-                            rdbReportCorruptRDB("Hash zipmap with dup elements");
+                        sds field = sdstrynewlen(fstr, flen);
+                        if (!field || dictAdd(dupSearchDict, field, NULL) != DICT_OK) {
+                            rdbReportCorruptRDB("Hash zipmap with dup elements, or big length (%u)", flen);
                             dictRelease(dupSearchDict);
                             sdsfree(field);
                             zfree(encoded);
diff --git a/src/sds.c b/src/sds.c
index f72eac921..6a85eb4aa 100644
--- a/src/sds.c
+++ b/src/sds.c
@@ -100,7 +100,7 @@ static inline size_t sdsTypeMaxSize(char type) {
  * You can print the string with printf() as there is an implicit \0 at the
  * end of the string. However the string is binary safe and can contain
  * \0 characters in the middle, as the length is stored in the sds header. */
-sds sdsnewlen(const void *init, size_t initlen) {
+sds _sdsnewlen(const void *init, size_t initlen, int trymalloc) {
     void *sh;
     sds s;
     char type = sdsReqType(initlen);
@@ -111,7 +111,9 @@ sds sdsnewlen(const void *init, size_t initlen) {
     unsigned char *fp; /* flags pointer. */
     size_t usable;
 
-    sh = s_malloc_usable(hdrlen+initlen+1, &usable);
+    sh = trymalloc?
+        s_trymalloc_usable(hdrlen+initlen+1, &usable) :
+        s_malloc_usable(hdrlen+initlen+1, &usable);
     if (sh == NULL) return NULL;
     if (init==SDS_NOINIT)
         init = NULL;
@@ -162,6 +164,14 @@ sds sdsnewlen(const void *init, size_t initlen) {
     return s;
 }
 
+sds sdsnewlen(const void *init, size_t initlen) {
+    return _sdsnewlen(init, initlen, 0);
+}
+
+sds sdstrynewlen(const void *init, size_t initlen) {
+    return _sdsnewlen(init, initlen, 1);
+}
+
 /* Create an empty (zero length) sds string. Even in this case the string
  * always has an implicit null term. */
 sds sdsempty(void) {
diff --git a/src/sds.h b/src/sds.h
index adcc12c0a..3a9e4cefe 100644
--- a/src/sds.h
+++ b/src/sds.h
@@ -216,6 +216,7 @@ static inline void sdssetalloc(sds s, size_t newlen) {
 }
 
 sds sdsnewlen(const void *init, size_t initlen);
+sds sdstrynewlen(const void *init, size_t initlen);
 sds sdsnew(const char *init);
 sds sdsempty(void);
 sds sdsdup(const sds s);
diff --git a/src/sdsalloc.h b/src/sdsalloc.h
index 725ad79e3..a1c5584f0 100644
--- a/src/sdsalloc.h
+++ b/src/sdsalloc.h
@@ -42,9 +42,13 @@
 #include "zmalloc.h"
 #define s_malloc zmalloc
 #define s_realloc zrealloc
+#define s_trymalloc ztrymalloc
+#define s_tryrealloc ztryrealloc
 #define s_free zfree
 #define s_malloc_usable zmalloc_usable
 #define s_realloc_usable zrealloc_usable
+#define s_trymalloc_usable ztrymalloc_usable
+#define s_tryrealloc_usable ztryrealloc_usable
 #define s_free_usable zfree_usable
 
 #endif
diff --git a/src/zmalloc.c b/src/zmalloc.c
index 7425198fe..a04ecacae 100644
--- a/src/zmalloc.c
+++ b/src/zmalloc.c
@@ -85,33 +85,44 @@ static void zmalloc_default_oom(size_t size) {
 
 static void (*zmalloc_oom_handler)(size_t) = zmalloc_default_oom;
 
-void *zmalloc(size_t size) {
+/* Try allocating memory, and return NULL if failed.
+ * '*usable' is set to the usable size if non NULL. */
+void *ztrymalloc_usable(size_t size, size_t *usable) {
     void *ptr = malloc(size+PREFIX_SIZE);
 
-    if (!ptr) zmalloc_oom_handler(size);
+    if (!ptr) return NULL;
 #ifdef HAVE_MALLOC_SIZE
-    update_zmalloc_stat_alloc(zmalloc_size(ptr));
+    size = zmalloc_size(ptr);
+    update_zmalloc_stat_alloc(size);
+    if (usable) *usable = size;
     return ptr;
 #else
     *((size_t*)ptr) = size;
     update_zmalloc_stat_alloc(size+PREFIX_SIZE);
+    if (usable) *usable = size;
     return (char*)ptr+PREFIX_SIZE;
 #endif
 }
 
-/* Similar to zmalloc, '*usable' is set to the usable size. */
-void *zmalloc_usable(size_t size, size_t *usable) {
-    void *ptr = malloc(size+PREFIX_SIZE);
+/* Allocate memory or panic */
+void *zmalloc(size_t size) {
+    void *ptr = ztrymalloc_usable(size, NULL);
+    if (!ptr) zmalloc_oom_handler(size);
+    return ptr;
+}
 
+/* Try allocating memory, and return NULL if failed. */
+void *ztrymalloc(size_t size) {
+    void *ptr = ztrymalloc_usable(size, NULL);
+    return ptr;
+}
+
+/* Allocate memory or panic.
+ * '*usable' is set to the usable size if non NULL. */
+void *zmalloc_usable(size_t size, size_t *usable) {
+    void *ptr = ztrymalloc_usable(size, usable);
     if (!ptr) zmalloc_oom_handler(size);
-#ifdef HAVE_MALLOC_SIZE
-    update_zmalloc_stat_alloc(*usable = zmalloc_size(ptr));
     return ptr;
-#else
-    *((size_t*)ptr) = *usable = size;
-    update_zmalloc_stat_alloc(size+PREFIX_SIZE);
-    return (char*)ptr+PREFIX_SIZE;
-#endif
 }
 
 /* Allocation and free functions that bypass the thread cache
@@ -132,101 +143,114 @@ void zfree_no_tcache(void *ptr) {
 }
 #endif
 
-void *zcalloc(size_t size) {
+/* Try allocating memory and zero it, and return NULL if failed.
+ * '*usable' is set to the usable size if non NULL. */
+void *ztrycalloc_usable(size_t size, size_t *usable) {
     void *ptr = calloc(1, size+PREFIX_SIZE);
+    if (ptr == NULL) return NULL;
 
-    if (!ptr) zmalloc_oom_handler(size);
 #ifdef HAVE_MALLOC_SIZE
-    update_zmalloc_stat_alloc(zmalloc_size(ptr));
+    size = zmalloc_size(ptr);
+    update_zmalloc_stat_alloc(size);
+    if (usable) *usable = size;
     return ptr;
 #else
     *((size_t*)ptr) = size;
     update_zmalloc_stat_alloc(size+PREFIX_SIZE);
+    if (usable) *usable = size;
     return (char*)ptr+PREFIX_SIZE;
 #endif
 }
 
-/* Similar to zcalloc, '*usable' is set to the usable size. */
-void *zcalloc_usable(size_t size, size_t *usable) {
-    void *ptr = calloc(1, size+PREFIX_SIZE);
+/* Allocate memory and zero it or panic */
+void *zcalloc(size_t size) {
+    void *ptr = ztrycalloc_usable(size, NULL);
+    if (!ptr) zmalloc_oom_handler(size);
+    return ptr;
+}
+
+/* Try allocating memory, and return NULL if failed. */
+void *ztrycalloc(size_t size) {
+    void *ptr = ztrycalloc_usable(size, NULL);
+    return ptr;
+}
 
+/* Allocate memory or panic.
+ * '*usable' is set to the usable size if non NULL. */
+void *zcalloc_usable(size_t size, size_t *usable) {
+    void *ptr = ztrycalloc_usable(size, usable);
     if (!ptr) zmalloc_oom_handler(size);
-#ifdef HAVE_MALLOC_SIZE
-    update_zmalloc_stat_alloc(*usable = zmalloc_size(ptr));
     return ptr;
-#else
-    *((size_t*)ptr) = *usable = size;
-    update_zmalloc_stat_alloc(size+PREFIX_SIZE);
-    return (char*)ptr+PREFIX_SIZE;
-#endif
 }
 
-void *zrealloc(void *ptr, size_t size) {
+/* Try reallocating memory, and return NULL if failed.
+ * '*usable' is set to the usable size if non NULL. */
+void *ztryrealloc_usable(void *ptr, size_t size, size_t *usable) {
 #ifndef HAVE_MALLOC_SIZE
     void *realptr;
 #endif
     size_t oldsize;
     void *newptr;
 
+    /* not allocating anything, just redirect to free. */
     if (size == 0 && ptr != NULL) {
         zfree(ptr);
+        if (usable) *usable = 0;
         return NULL;
     }
-    if (ptr == NULL) return zmalloc(size);
+    /* Not freeing anything, just redirect to malloc. */
+    if (ptr == NULL)
+        return ztrymalloc_usable(size, usable);
+
 #ifdef HAVE_MALLOC_SIZE
     oldsize = zmalloc_size(ptr);
     newptr = realloc(ptr,size);
-    if (!newptr) zmalloc_oom_handler(size);
+    if (newptr == NULL) {
+        if (usable) *usable = 0;
+        return NULL;
+    }
 
     update_zmalloc_stat_free(oldsize);
-    update_zmalloc_stat_alloc(zmalloc_size(newptr));
+    size = zmalloc_size(newptr);
+    update_zmalloc_stat_alloc(size);
+    if (usable) *usable = size;
     return newptr;
 #else
     realptr = (char*)ptr-PREFIX_SIZE;
     oldsize = *((size_t*)realptr);
     newptr = realloc(realptr,size+PREFIX_SIZE);
-    if (!newptr) zmalloc_oom_handler(size);
+    if (newptr == NULL) {
+        if (usable) *usable = 0;
+        return NULL;
+    }
 
     *((size_t*)newptr) = size;
-    update_zmalloc_stat_free(oldsize+PREFIX_SIZE);
-    update_zmalloc_stat_alloc(size+PREFIX_SIZE);
+    update_zmalloc_stat_free(oldsize);
+    update_zmalloc_stat_alloc(size);
+    if (usable) *usable = size;
     return (char*)newptr+PREFIX_SIZE;
 #endif
 }
 
-/* Similar to zrealloc, '*usable' is set to the new usable size. */
-void *zrealloc_usable(void *ptr, size_t size, size_t *usable) {
-#ifndef HAVE_MALLOC_SIZE
-    void *realptr;
-#endif
-    size_t oldsize;
-    void *newptr;
-
-    if (size == 0 && ptr != NULL) {
-        zfree(ptr);
-        *usable = 0;
-        return NULL;
-    }
-    if (ptr == NULL) return zmalloc_usable(size, usable);
-#ifdef HAVE_MALLOC_SIZE
-    oldsize = zmalloc_size(ptr);
-    newptr = realloc(ptr,size);
-    if (!newptr) zmalloc_oom_handler(size);
+/* Reallocate memory and zero it or panic */
+void *zrealloc(void *ptr, size_t size) {
+    ptr = ztryrealloc_usable(ptr, size, NULL);
+    if (!ptr && size != 0) zmalloc_oom_handler(size);
+    return ptr;
+}
 
-    update_zmalloc_stat_free(oldsize);
-    update_zmalloc_stat_alloc(*usable = zmalloc_size(newptr));
-    return newptr;
-#else
-    realptr = (char*)ptr-PREFIX_SIZE;
-    oldsize = *((size_t*)realptr);
-    newptr = realloc(realptr,size+PREFIX_SIZE);
-    if (!newptr) zmalloc_oom_handler(size);
+/* Try Reallocating memory, and return NULL if failed. */
+void *ztryrealloc(void *ptr, size_t size) {
+    ptr = ztryrealloc_usable(ptr, size, NULL);
+    return ptr;
+}
 
-    *((size_t*)newptr) = *usable = size;
-    update_zmalloc_stat_free(oldsize);
-    update_zmalloc_stat_alloc(size);
-    return (char*)newptr+PREFIX_SIZE;
-#endif
+/* Reallocate memory or panic.
+ * '*usable' is set to the usable size if non NULL. */
+void *zrealloc_usable(void *ptr, size_t size, size_t *usable) {
+    ptr = ztryrealloc_usable(ptr, size, usable);
+    if (!ptr && size != 0) zmalloc_oom_handler(size);
+    return ptr;
 }
 
 /* Provide zmalloc_size() for systems where this function is not provided by
diff --git a/src/zmalloc.h b/src/zmalloc.h
index df63253ef..64bc9fc76 100644
--- a/src/zmalloc.h
+++ b/src/zmalloc.h
@@ -80,10 +80,16 @@
 void *zmalloc(size_t size);
 void *zcalloc(size_t size);
 void *zrealloc(void *ptr, size_t size);
+void *ztrymalloc(size_t size);
+void *ztrycalloc(size_t size);
+void *ztryrealloc(void *ptr, size_t size);
 void zfree(void *ptr);
 void *zmalloc_usable(size_t size, size_t *usable);
 void *zcalloc_usable(size_t size, size_t *usable);
 void *zrealloc_usable(void *ptr, size_t size, size_t *usable);
+void *ztrymalloc_usable(size_t size, size_t *usable);
+void *ztrycalloc_usable(size_t size, size_t *usable);
+void *ztryrealloc_usable(void *ptr, size_t size, size_t *usable);
 void zfree_usable(void *ptr, size_t *usable);
 char *zstrdup(const char *s);
 size_t zmalloc_used_memory(void);
diff --git a/tests/integration/corrupt-dump.tcl b/tests/integration/corrupt-dump.tcl
index fbcca2f2a..4b09d45b9 100644
--- a/tests/integration/corrupt-dump.tcl
+++ b/tests/integration/corrupt-dump.tcl
@@ -444,6 +444,25 @@ test {corrupt payload: fuzzer findings - hash convert asserts on RESTORE with sh
     }
 }
 
+test {corrupt payload: OOM in rdbGenericLoadStringObject} {
+    start_server [list overrides [list loglevel verbose use-exit-on-panic yes crash-memcheck-enabled no] ] {
+        r config set sanitize-dump-payload no
+        catch { r RESTORE x 0 "\x0A\x81\x7F\xFF\xFF\xFF\xFF\xFF\xFF\xFF\x13\x00\x00\x00\x0E\x00\x00\x00\x02\x00\x00\x02\x61\x00\x04\x02\x62\x00\xFF\x09\x00\x57\x04\xE5\xCD\xD4\x37\x6C\x57" } err
+        assert_match "*Bad data format*" $err
+        r ping
+    }
+}
+
+test {corrupt payload: fuzzer findings - OOM in dictExpand} {
+    start_server [list overrides [list loglevel verbose use-exit-on-panic yes crash-memcheck-enabled no] ] {
+        r config set sanitize-dump-payload no
+        r debug set-skip-checksum-validation 1
+        catch { r RESTORE x 0 "\x02\x81\x02\x5F\x31\xC0\x00\xC0\x02\x09\x00\xCD\x84\x2C\xB7\xE8\xA4\x49\x57" } err
+        assert_match "*Bad data format*" $err
+        r ping
+    }
+}
+
 test {corrupt payload: fuzzer findings - invalid tail offset after removal} {
     start_server [list overrides [list loglevel verbose use-exit-on-panic yes crash-memcheck-enabled no] ] {
         r config set sanitize-dump-payload no