From a3f1535b5709bba940f62078b6ce4e9c96ab4f87 Mon Sep 17 00:00:00 2001
From: Madelyn Olson <madelyneolson@gmail.com>
Date: Mon, 10 Jun 2024 12:30:57 -0700
Subject: [PATCH] Fix misuse of safe iterators (#612)

Safe iterators must call resetIterators when they are done being used.
Fix one issue where a safe iterator was not correctly calling reset
during cluster slot caching and fixed a second issue where reset
iterator was being called twice. For the double release case,
kvstoreIteratorNextDict is responsible for patching up the iterator, but
we were calling it a second time in kvstoreIteratorNext.

In addition, I added some documentation around initializing iterators,
added an assert to prevent double initialization, and remove a function
from the public interface which isn't needed and might lead to incorrect
usage of the safe iterators.

Bumping srgsanky for finding it here:
https://github.com/valkey-io/valkey/commit/c4782066e76ba20848fab95fd4909507cc069c3b#r142867004.

---------

Signed-off-by: Madelyn Olson <madelyneolson@gmail.com>
---
 src/cluster_legacy.c |  4 ++--
 src/dict.c           |  9 +++++++--
 src/kvstore.c        | 11 +++--------
 src/kvstore.h        |  1 -
 4 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/src/cluster_legacy.c b/src/cluster_legacy.c
index ef78f29434..cd3786fe05 100644
--- a/src/cluster_legacy.c
+++ b/src/cluster_legacy.c
@@ -4888,7 +4888,7 @@ void bitmapClearBit(unsigned char *bitmap, int pos) {
  * MIGRATE_TO flag the when a primary gets the first slot. */
 int clusterPrimariesHaveReplicas(void) {
     dictIterator di;
-    dictInitSafeIterator(&di, server.cluster->nodes);
+    dictInitIterator(&di, server.cluster->nodes);
     dictEntry *de;
     int replicas = 0;
     while ((de = dictNext(&di)) != NULL) {
@@ -6509,7 +6509,7 @@ void clusterPromoteSelfToPrimary(void) {
 
 int detectAndUpdateCachedNodeHealth(void) {
     dictIterator di;
-    dictInitSafeIterator(&di, server.cluster->nodes);
+    dictInitIterator(&di, server.cluster->nodes);
     dictEntry *de;
     clusterNode *node;
     int overall_health_changed = 0;
diff --git a/src/dict.c b/src/dict.c
index dc227dbee9..bc92d49564 100644
--- a/src/dict.c
+++ b/src/dict.c
@@ -943,6 +943,8 @@ unsigned long long dictFingerprint(dict *d) {
     return hash;
 }
 
+/* Initiaize a normal iterator. This function should be called when initializing
+ * an iterator on the stack. */
 void dictInitIterator(dictIterator *iter, dict *d) {
     iter->d = d;
     iter->table = 0;
@@ -952,6 +954,8 @@ void dictInitIterator(dictIterator *iter, dict *d) {
     iter->nextEntry = NULL;
 }
 
+/* Initialize a safe iterator, which is allowed to modify the dictionary while iterating.
+ * You must call dictResetIterator when you are done with a safe iterator. */
 void dictInitSafeIterator(dictIterator *iter, dict *d) {
     dictInitIterator(iter, d);
     iter->safe = 1;
@@ -959,9 +963,10 @@ void dictInitSafeIterator(dictIterator *iter, dict *d) {
 
 void dictResetIterator(dictIterator *iter) {
     if (!(iter->index == -1 && iter->table == 0)) {
-        if (iter->safe)
+        if (iter->safe) {
             dictResumeRehashing(iter->d);
-        else
+            assert(iter->d->pauserehash >= 0);
+        } else
             assert(iter->fingerprint == dictFingerprint(iter->d));
     }
 }
diff --git a/src/kvstore.c b/src/kvstore.c
index 8de74c724b..a43b72e1e1 100644
--- a/src/kvstore.c
+++ b/src/kvstore.c
@@ -48,6 +48,8 @@
 
 #define UNUSED(V) ((void)V)
 
+static dict *kvstoreIteratorNextDict(kvstoreIterator *kvs_it);
+
 struct _kvstore {
     int flags;
     dictType dtype;
@@ -572,7 +574,7 @@ void kvstoreIteratorRelease(kvstoreIterator *kvs_it) {
 }
 
 /* Returns next dictionary from the iterator, or NULL if iteration is complete. */
-dict *kvstoreIteratorNextDict(kvstoreIterator *kvs_it) {
+static dict *kvstoreIteratorNextDict(kvstoreIterator *kvs_it) {
     if (kvs_it->next_didx == -1) return NULL;
 
     /* The dict may be deleted during the iteration process, so here need to check for NULL. */
@@ -600,13 +602,6 @@ dictEntry *kvstoreIteratorNext(kvstoreIterator *kvs_it) {
     if (!de) { /* No current dict or reached the end of the dictionary. */
         dict *d = kvstoreIteratorNextDict(kvs_it);
         if (!d) return NULL;
-        if (kvs_it->di.d) {
-            /* Before we move to the next dict, reset the iter of the previous dict. */
-            dictIterator *iter = &kvs_it->di;
-            dictResetIterator(iter);
-            /* In the safe iterator context, we may delete entries. */
-            freeDictIfNeeded(kvs_it->kvs, kvs_it->didx);
-        }
         dictInitSafeIterator(&kvs_it->di, d);
         de = dictNext(&kvs_it->di);
     }
diff --git a/src/kvstore.h b/src/kvstore.h
index d3c5949d1f..e7e21f8aa9 100644
--- a/src/kvstore.h
+++ b/src/kvstore.h
@@ -40,7 +40,6 @@ uint64_t kvstoreGetHash(kvstore *kvs, const void *key);
 /* kvstore iterator specific functions */
 kvstoreIterator *kvstoreIteratorInit(kvstore *kvs);
 void kvstoreIteratorRelease(kvstoreIterator *kvs_it);
-dict *kvstoreIteratorNextDict(kvstoreIterator *kvs_it);
 int kvstoreIteratorGetCurrentDictIndex(kvstoreIterator *kvs_it);
 dictEntry *kvstoreIteratorNext(kvstoreIterator *kvs_it);