From 0049bc03ceae07155626a6e47482e26001b9fcd1 Mon Sep 17 00:00:00 2001
From: James Park-Watt <jimmypw@users.noreply.github.com>
Date: Sun, 27 Oct 2024 11:33:53 +0000
Subject: [PATCH] Add periodic in-memory store pruner (#200)

---
 pkg/storage/memory/memory.go | 38 ++++++++++++++++++++++++++++++++----
 1 file changed, 34 insertions(+), 4 deletions(-)

diff --git a/pkg/storage/memory/memory.go b/pkg/storage/memory/memory.go
index d533aab..694e841 100644
--- a/pkg/storage/memory/memory.go
+++ b/pkg/storage/memory/memory.go
@@ -18,14 +18,37 @@ type (
 
 	storageMem struct {
 		sync.RWMutex
-		store map[string]memStorageSecret
+		store           map[string]memStorageSecret
+		storePruneTimer *time.Ticker
 	}
 )
 
 // New creates a new In-Mem storage
 func New() storage.Storage {
-	return &storageMem{
-		store: make(map[string]memStorageSecret),
+	store := &storageMem{
+		store:           make(map[string]memStorageSecret),
+		storePruneTimer: time.NewTicker(time.Minute),
+	}
+
+	go store.storePruner()
+
+	return store
+}
+
+func (s *storageMem) storePruner() {
+	for range s.storePruneTimer.C {
+		s.pruneStore()
+	}
+}
+
+func (s *storageMem) pruneStore() {
+	s.Lock()
+	defer s.Unlock()
+
+	for k, v := range s.store {
+		if v.hasExpired() {
+			delete(s.store, k)
+		}
 	}
 }
 
@@ -68,9 +91,16 @@ func (s *storageMem) ReadAndDestroy(id string) (string, error) {
 
 	defer delete(s.store, id)
 
-	if !secret.Expiry.IsZero() && secret.Expiry.Before(time.Now()) {
+	// Still check to see if the secret has expired in order to prevent a
+	// race condition where a secret has expired but the the store pruner has
+	// not yet been invoked.
+	if secret.hasExpired() {
 		return "", storage.ErrSecretNotFound
 	}
 
 	return secret.Secret, nil
 }
+
+func (m *memStorageSecret) hasExpired() bool {
+	return !m.Expiry.IsZero() && m.Expiry.Before(time.Now())
+}