From 79b37ff53548ee4d0fca287cb37a2eaf84c519a7 Mon Sep 17 00:00:00 2001
From: Moti Cohen <moticless@gmail.com>
Date: Wed, 28 May 2025 08:02:10 +0300
Subject: [PATCH] Fix RESTORE with TTL (#14071)

restoreCommand() creates a key-value object (kv) with a TTL in two steps.
During the second step, setExpire() may reallocate the kv object. To ensure
correct behavior, kv must be updated after this call, as it might be used later
in the function.
---
 src/cluster.c       |  2 +-
 src/db.c            |  8 +++++---
 src/expire.c        |  2 +-
 src/module.c        |  3 ++-
 src/t_string.c      |  2 +-
 tests/unit/dump.tcl | 14 ++++++++++++++
 6 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/src/cluster.c b/src/cluster.c
index 070c278b9..e3f6e90d7 100644
--- a/src/cluster.c
+++ b/src/cluster.c
@@ -250,7 +250,7 @@ void restoreCommand(client *c) {
     }
 
     if (ttl) {
-        setExpire(c,c->db,key,ttl);
+        kv = setExpire(c,c->db,key,ttl); /* might realloc kvobj */
         if (!absttl) {
             /* Propagate TTL as absolute timestamp */
             robj *ttl_obj = createStringObjectFromLongLong(ttl);
diff --git a/src/db.c b/src/db.c
index 06eefc0cb..32184f580 100644
--- a/src/db.c
+++ b/src/db.c
@@ -1818,7 +1818,7 @@ void renameGenericCommand(client *c, int nx) {
 
     dbDelete(c->db,c->argv[1]);
     dbAdd(c->db, c->argv[2], &o);
-    if (expire != -1) setExpire(c, c->db, c->argv[2], expire);
+    if (expire != -1) o = setExpire(c, c->db, c->argv[2], expire);
 
     /* If hash with HFEs, register in db->hexpires */
     if (minHashExpireTime != EB_EXPIRE_TIME_INVALID)
@@ -2021,7 +2021,7 @@ void copyCommand(client *c) {
 
     /* if key with expiration then set it */
     if (expire != -1)
-        setExpire(c, dst, newkey, expire);
+        newobj = setExpire(c, dst, newkey, expire);
 
     /* If minExpiredField was set, then the object is hash with expiration
      * on fields and need to register it in global HFE DS */
@@ -2240,7 +2240,9 @@ int removeExpire(redisDb *db, robj *key) {
 /* Set an expire to the specified key. If the expire is set in the context
  * of an user calling a command 'c' is the client, otherwise 'c' is set
  * to NULL. The 'when' parameter is the absolute unix time in milliseconds
- * after which the key will no longer be considered valid. */
+ * after which the key will no longer be considered valid.
+ * 
+ * Note: It may reallocate kvobj. The returned ref may point to a new object. */
 kvobj *setExpire(client *c, redisDb *db, robj *key, long long when) {
     return setExpireByLink(c,db,key->ptr,when,NULL);
 }
diff --git a/src/expire.c b/src/expire.c
index 1b3492b09..d9a417ecc 100644
--- a/src/expire.c
+++ b/src/expire.c
@@ -719,7 +719,7 @@ void expireGenericCommand(client *c, long long basetime, int unit) {
         addReply(c, shared.cone);
         return;
     } else {
-        setExpire(c,c->db,key,when);
+        kv = setExpire(c,c->db,key,when); /* might realloc kv */
         addReply(c,shared.cone);
         /* Propagate as PEXPIREAT millisecond-timestamp
          * Only rewrite the command arg if not already PEXPIREAT */
diff --git a/src/module.c b/src/module.c
index 496d606d4..9c795cb18 100644
--- a/src/module.c
+++ b/src/module.c
@@ -4269,7 +4269,8 @@ int RM_SetExpire(RedisModuleKey *key, mstime_t expire) {
         return REDISMODULE_ERR;
     if (expire != REDISMODULE_NO_EXPIRE) {
         expire += commandTimeSnapshot();
-        setExpire(key->ctx->client,key->db,key->key,expire);
+        /* setExpire() might realloc kvobj */ 
+        key->kv = setExpire(key->ctx->client,key->db,key->key,expire);
     } else {
         removeExpire(key->db,key->key);
     }
diff --git a/src/t_string.c b/src/t_string.c
index 64444c778..34806541c 100644
--- a/src/t_string.c
+++ b/src/t_string.c
@@ -397,7 +397,7 @@ void getexCommand(client *c) {
         notifyKeyspaceEvent(NOTIFY_GENERIC, "del", c->argv[1], c->db->id);
         server.dirty++;
     } else if (expire) {
-        setExpire(c,c->db,c->argv[1],milliseconds);
+        o = setExpire(c,c->db,c->argv[1],milliseconds);
         /* Propagate as PXEXPIREAT millisecond-timestamp if there is
          * EX/PX/EXAT/PXAT flag and the key has not expired. */
         robj *milliseconds_obj = createStringObjectFromLongLong(milliseconds);
diff --git a/tests/unit/dump.tcl b/tests/unit/dump.tcl
index 0cd62804a..a03c846dd 100644
--- a/tests/unit/dump.tcl
+++ b/tests/unit/dump.tcl
@@ -59,6 +59,20 @@ start_server {tags {"dump"}} {
         assert_equal [r get foo] {bar}
         r config set maxmemory-policy noeviction
     } {OK} {needs:config-maxmemory}
+    
+    test {RESTORE with TTL maintain valid object} {
+        # RESTORE Creates a string with TTL in two steps. The second step potentially 
+        # reallocates the object. Access the object and verify it is not corrupted
+        r del foo
+        r set foo bar
+        set encoded [r dump foo]
+        # Iterate several times and verify it is consistent
+        for {set i 0} {$i < 100} {incr i} {
+            r del foo
+            r restore foo 1000 $encoded IDLETIME 500
+            assert_equal [r get foo] {bar}
+        }
+    }
 
     test {RESTORE can set LFU} {
         r set foo bar