From b7934083d0c04e213475bb8c719e1e3b65201f37 Mon Sep 17 00:00:00 2001
From: James Park-Watt <james@park-watt.co.uk>
Date: Sat, 26 Oct 2024 21:43:31 +0100
Subject: [PATCH] feat: periodic in-memory store pruner

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

diff --git a/pkg/storage/memory/memory.go b/pkg/storage/memory/memory.go
index d533aab..0fd2b2b 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
+		storePruneTimer time.Ticker
+		store           map[string]memStorageSecret
 	}
 )
 
 // New creates a new In-Mem storage
 func New() storage.Storage {
-	return &storageMem{
-		store: make(map[string]memStorageSecret),
+	store := &storageMem{
+		storePruneTimer: *time.NewTicker(time.Minute),
+		store:           make(map[string]memStorageSecret),
+	}
+
+	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,19 @@ 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 {
+	if !m.Expiry.IsZero() && m.Expiry.Before(time.Now()) {
+		return true
+	}
+	return false
+}