From dfa37fc7faf44eb5440fbed5e1a8eed8d5c7bcfc Mon Sep 17 00:00:00 2001
From: 1e99 <1e99@1e99.eu>
Date: Wed, 30 Oct 2024 00:21:25 +0100
Subject: [PATCH] allow storage to close

---
 main.go            |  2 ++
 storage/ram.go     | 43 +++++++++++++++++++++++++------------------
 storage/storage.go |  1 +
 3 files changed, 28 insertions(+), 18 deletions(-)

diff --git a/main.go b/main.go
index d994d52..913472a 100644
--- a/main.go
+++ b/main.go
@@ -20,6 +20,8 @@ func run() error {
 		return err
 	}
 
+	defer storage.Close()
+
 	mux := http.NewServeMux()
 	mux.Handle("GET /", routes.ServeFiles(embedFS, "static"))
 	mux.Handle("POST /api/password", routes.CreatePassword(storage, 12*1024, base64.StdEncoding))
diff --git a/storage/ram.go b/storage/ram.go
index 84e8de6..6ff4dc9 100644
--- a/storage/ram.go
+++ b/storage/ram.go
@@ -10,23 +10,17 @@ func NewRamStore() Store {
 	store := &ram{
 		passwords: make(map[string]entry),
 		lock:      sync.Mutex{},
+		close:     make(chan bool),
 	}
 
-	go func() {
-		ticker := time.Tick(25 * time.Second)
-
-		for {
-			store.clearExpired()
-			<-ticker
-		}
-	}()
-
+	go store.clearExpired()
 	return store
 }
 
 type ram struct {
 	passwords map[string]entry
 	lock      sync.Mutex
+	close     chan bool
 }
 
 func (store *ram) CreatePassword(password []byte, expiresAt time.Time) (string, error) {
@@ -64,18 +58,31 @@ func (store *ram) GetPassword(id string) ([]byte, error) {
 	return password.password, nil
 }
 
-func (store *ram) clearExpired() error {
-	store.lock.Lock()
-	defer store.lock.Unlock()
-	time := time.Now()
+func (store *ram) Close() error {
+	store.close <- true
+	return nil
+}
 
-	for id, password := range store.passwords {
-		if time.After(password.expiresAt) {
-			delete(store.passwords, id)
+func (store *ram) clearExpired() error {
+	ticker := time.NewTicker(20 * time.Second)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-store.close:
+			return nil
+		case <-ticker.C:
+			store.lock.Lock()
+			defer store.lock.Unlock()
+			time := time.Now()
+
+			for id, password := range store.passwords {
+				if time.After(password.expiresAt) {
+					delete(store.passwords, id)
+				}
+			}
 		}
 	}
-
-	return nil
 }
 
 func (store *ram) generateId(length int) string {
diff --git a/storage/storage.go b/storage/storage.go
index 4571e87..8151500 100644
--- a/storage/storage.go
+++ b/storage/storage.go
@@ -21,6 +21,7 @@ type entry struct {
 type Store interface {
 	CreatePassword(password []byte, expiresAt time.Time) (string, error)
 	GetPassword(id string) ([]byte, error)
+	Close() error
 }
 
 func NewStore() (Store, error) {