Fix expired-key enumeration

Core/MMKV.h

@@ -577,7 +577,8 @@ public:
     bool removeValuesForKeys(NSArray *arrKeys);
 
     typedef void (^EnumerateBlock)(NSString *key, BOOL *stop);
-    void enumerateKeys(EnumerateBlock block);
+    // filterExpire: filter expired keys at the start, keep in mind it comes with additional cost.
+    void enumerateKeys(EnumerateBlock block, bool filterExpire = false);
 #endif // __OBJC__
     // filterExpire: return all non-expired keys, keep in mind it comes with cost
     std::vector<std::string> allKeys(bool filterExpire = false);

Core/MMKV_OSX.cpp

@@ -560,7 +560,7 @@ bool MMKV::removeValuesForKeys(const std::vector<std::string> &arrKeys) {
     return ret;
 }
 
-void MMKV::enumerateKeys(EnumerateBlock block) {
+void MMKV::enumerateKeys(EnumerateBlock block, bool filterExpire) {
     if (block == nil) {
         return;
     }
@@ -568,7 +568,20 @@ void MMKV::enumerateKeys(EnumerateBlock block) {
     checkLoadData();
 
     MMKVInfo("enumerate [%s] begin", m_mmapID.c_str());
-    if (m_crypter) {
+    if (m_enableKeyExpire) {
+        // A callback can remove expired entries just by reading them. Snapshot
+        // and retain the keys so both iteration and callback key lifetimes stay valid.
+        @autoreleasepool {
+            auto keys = allKeysObjC(filterExpire);
+            for (NSString *key in keys) {
+                BOOL stop = NO;
+                block(key, &stop);
+                if (stop) {
+                    break;
+                }
+            }
+        }
+    } else if (m_crypter) {
         for (const auto &pair : *m_dicCrypt) {
             BOOL stop = NO;
             block(pair.first, &stop);

iOS/MMKV/MMKV/MMKV.h

@@ -359,7 +359,10 @@ class MMKV;
 
 + (NSString *)version;
 
+/// Includes stored expired keys; callbacks can safely read their values.
 - (void)enumerateKeys:(void (^)(NSString *key, BOOL *stop))block;
+/// Enumerates non-expired keys at the start, with additional cost.
+- (void)enumerateNonExpiredKeys:(void (^)(NSString *key, BOOL *stop))block;
 - (NSArray *)allKeys;
 
 /// return count of non-expired keys, keep in mind that it comes with cost

iOS/MMKV/MMKV/libMMKV.mm

@@ -796,6 +796,10 @@ static BOOL g_hasCalledInitializeMMKV = NO;
     m_mmkv->enumerateKeys(block);
 }
 
+- (void)enumerateNonExpiredKeys:(void (^)(NSString *key, BOOL *stop))block {
+    m_mmkv->enumerateKeys(block, true);
+}
+
 - (NSArray *)allKeys {
     return m_mmkv->allKeysObjC();
 }

iOS/MMKVDemo/MMKVDemoTests/MMKVDemoTests.mm

@@ -427,6 +427,109 @@ using namespace std;
     XCTAssertEqualObjects(sValue, @"hello");
 }
 
+- (void)checkEnumerationWithCryptKey:(NSData *)cryptKey
+            expireDuringEnumeration:(BOOL)expireDuringEnumeration
+                     nonExpiredOnly:(BOOL)nonExpiredOnly {
+    NSString *mmapID = [@"enumeration_expiration_" stringByAppendingString:NSUUID.UUID.UUIDString];
+    MMKV *kv = [MMKV mmkvWithID:mmapID cryptKey:cryptKey];
+    XCTAssertNotNil(kv);
+    XCTAssertTrue([kv enableAutoKeyExpire:MMKVExpireNever]);
+
+    __block NSUInteger emptyCount = 0;
+    [kv enumerateNonExpiredKeys:^(NSString *key, BOOL *stop) {
+        ++emptyCount;
+    }];
+    XCTAssertEqual(emptyCount, (NSUInteger) 0);
+
+    // Leave the dictionary as the sole owner of these non-tagged NSStrings.
+    @autoreleasepool {
+        for (NSUInteger index = 0; index < 16; ++index) {
+            NSString *key = [NSString stringWithFormat:@"expired-dynamically-allocated-enumeration-key-%lu",
+                                                      (unsigned long) index];
+            XCTAssertTrue([kv setInt32:42 forKey:key expireDuration:expireDuringEnumeration ? 2 : 1]);
+        }
+    }
+    XCTAssertTrue([kv setInt32:7 forKey:@"live-key" expireDuration:MMKVExpireNever]);
+
+    if (!expireDuringEnumeration) {
+        [NSThread sleepForTimeInterval:1.1];
+    }
+    // enumerateKeys: includes expired entries without filtering.
+    XCTAssertEqual(kv.count, (size_t) 17);
+    __block NSUInteger unfilteredCount = 0;
+    [kv enumerateKeys:^(NSString *key, BOOL *stop) {
+        ++unfilteredCount;
+    }];
+    XCTAssertEqual(unfilteredCount, (NSUInteger) 17);
+    __block NSUInteger stoppedCount = 0;
+    [kv enumerateKeys:^(NSString *key, BOOL *stop) {
+        ++stoppedCount;
+        *stop = YES;
+    }];
+    XCTAssertEqual(stoppedCount, (NSUInteger) 1);
+
+    NSMutableSet<NSString *> *visited = [NSMutableSet set];
+    void (^readValue)(NSString *, BOOL *) = ^(NSString *key, BOOL *stop) {
+        if (expireDuringEnumeration && visited.count == 0) {
+            // Filtering once before iteration cannot protect against this case.
+            [NSThread sleepForTimeInterval:2.1];
+        }
+        BOOL expired = [key hasPrefix:@"expired-"];
+        if ([key hasSuffix:@"0"]) {
+            XCTAssertFalse([kv containsKey:key]);
+        } else {
+            BOOL hasValue = NO;
+            XCTAssertEqual([kv getInt32ForKey:key defaultValue:-1 hasValue:&hasValue], expired ? -1 : 7);
+            XCTAssertEqual(hasValue, !expired);
+        }
+        // The getter can erase the entry; the callback key must remain usable.
+        XCTAssertFalse([visited containsObject:key]);
+        [visited addObject:key];
+    };
+    if (nonExpiredOnly) {
+        [kv enumerateNonExpiredKeys:readValue];
+    } else {
+        [kv enumerateKeys:readValue];
+    }
+    XCTAssertEqual(visited.count, (NSUInteger) ((!nonExpiredOnly || expireDuringEnumeration) ? 17 : 1));
+    XCTAssertTrue([visited containsObject:@"live-key"]);
+    XCTAssertEqual(kv.count, (size_t) 1);
+    XCTAssertEqual([kv getInt32ForKey:@"live-key"], 7);
+
+    XCTAssertTrue([kv setInt32:8 forKey:@"another-live-key" expireDuration:MMKVExpireNever]);
+    stoppedCount = 0;
+    [kv enumerateNonExpiredKeys:^(NSString *key, BOOL *stop) {
+        ++stoppedCount;
+        *stop = YES;
+    }];
+    XCTAssertEqual(stoppedCount, (NSUInteger) 1);
+    [kv close];
+    XCTAssertTrue([MMKV removeStorage:mmapID rootPath:nil]);
+}
+
+- (void)checkEnumerationWithCryptKey:(NSData *)cryptKey expireDuringEnumeration:(BOOL)expireDuringEnumeration {
+    [self checkEnumerationWithCryptKey:cryptKey expireDuringEnumeration:expireDuringEnumeration nonExpiredOnly:NO];
+    [self checkEnumerationWithCryptKey:cryptKey expireDuringEnumeration:expireDuringEnumeration nonExpiredOnly:YES];
+}
+
+- (void)testEnumerateExpiredKeys {
+    [self checkEnumerationWithCryptKey:nil expireDuringEnumeration:NO];
+}
+
+- (void)testEnumerateExpiredKeysEncrypted {
+    [self checkEnumerationWithCryptKey:[@"enumerate-secret" dataUsingEncoding:NSUTF8StringEncoding]
+             expireDuringEnumeration:NO];
+}
+
+- (void)testKeysExpireDuringEnumeration {
+    [self checkEnumerationWithCryptKey:nil expireDuringEnumeration:YES];
+}
+
+- (void)testKeysExpireDuringEnumerationEncrypted {
+    [self checkEnumerationWithCryptKey:[@"enumerate-secret" dataUsingEncoding:NSUTF8StringEncoding]
+             expireDuringEnumeration:YES];
+}
+
 - (void)testExpirationOverflow {
     MMKV *kv = [MMKV mmkvWithID:@"expiration_overflow_test"];
     [kv clearAll];