From 22f4401661861dc978af7b7c824f115c23bb85a3 Mon Sep 17 00:00:00 2001
From: Juan Calderon-Perez <jgcalderonperez@protonmail.com>
Date: Sun, 29 Oct 2023 21:00:43 -0400
Subject: [PATCH] Add fix to verify Keys are expired during Keys() call.

---
 memory/memory.go      | 13 +++++++++++--
 memory/memory_test.go | 31 +++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/memory/memory.go b/memory/memory.go
index 20a5076c..cf35923c 100644
--- a/memory/memory.go
+++ b/memory/memory.go
@@ -151,9 +151,18 @@ func (s *Storage) Keys() ([][]byte, error) {
 		return nil, nil
 	}
 
+	ts := atomic.LoadUint32(&internal.Timestamp)
 	keys := make([][]byte, 0, len(s.db))
-	for key := range s.db {
-		keys = append(keys, []byte(key))
+	for key, v := range s.db {
+		// Filter out the expired keys
+		if v.expiry == 0 || v.expiry > ts {
+			keys = append(keys, []byte(key))
+		}
+	}
+
+	// Double check if no valid keys were found
+	if len(keys) == 0 {
+		return nil, nil
 	}
 
 	return keys, nil
diff --git a/memory/memory_test.go b/memory/memory_test.go
index 71d8bf3d..5cf4214a 100644
--- a/memory/memory_test.go
+++ b/memory/memory_test.go
@@ -81,6 +81,37 @@ func Test_Storage_Memory_Set_Expiration(t *testing.T) {
 	require.Nil(t, keys)
 }
 
+func Test_Storage_Memory_Set_Long_Expiration_with_Keys(t *testing.T) {
+	var (
+		testStore = New()
+		key       = "john"
+		val       = []byte("doe")
+		exp       = 5 * time.Second
+	)
+
+	keys, err := testStore.Keys()
+	require.NoError(t, err)
+	require.Nil(t, keys)
+
+	err = testStore.Set(key, val, exp)
+	require.NoError(t, err)
+
+	time.Sleep(1100 * time.Millisecond)
+
+	keys, err = testStore.Keys()
+	require.NoError(t, err)
+	require.Len(t, keys, 1)
+
+	time.Sleep(4000 * time.Millisecond)
+	result, err := testStore.Get(key)
+	require.NoError(t, err)
+	require.Zero(t, len(result))
+
+	keys, err = testStore.Keys()
+	require.NoError(t, err)
+	require.Nil(t, keys)
+}
+
 func Test_Storage_Memory_Get_NotExist(t *testing.T) {
 	testStore := New()
 	result, err := testStore.Get("notexist")