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) {