From 990fddff7ae05f5bd0be476fbbd8868db5a03a23 Mon Sep 17 00:00:00 2001
From: Knut Ahlers <knut@ahlers.me>
Date: Fri, 24 Jan 2020 16:33:14 +0100
Subject: [PATCH] [#12] Add lazy-expiry to mem-storage, unify envvars

Signed-off-by: Knut Ahlers <knut@ahlers.me>
---
 README.md        |  5 ++---
 storage_mem.go   | 45 +++++++++++++++++++++++++++++++++++++++------
 storage_redis.go | 13 ++++++++++---
 3 files changed, 51 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index b3dbebe..e85615f 100644
--- a/README.md
+++ b/README.md
@@ -21,12 +21,11 @@
 For a better setup you can choose the backend which is used to store the secrets:
 
 - `mem` - In memory storage (wiped on restart of the daemon)
+  - `SECRET_EXPIRY` - Expiry of the keys in seconds (Default `0` = no expiry)
 - `redis` - Storing the secrets in a hash under one key
   - `REDIS_URL` - Redis connection string `tcp://auth:PWD@HOST:PORT/DB`
-  - `REDIS_EXPIRY` - Expiry of the keys in seconds (Default `0` = no expiry)
   - `REDIS_KEY` - Key prefix to store the keys under (Default `io.luzifer.ots`)
-
-**Hint:** Starting in `v0.7.0` the secrets in Redis are no longer stored in a hash but in own keys. This allows for individual expiry. At the first start of `v0.7.0` the old data will be migrated automatically and afterwards be subject of expiry if you set `REDIS_EXPIRY`. My hosted instance uses an expiry of 90d (= 7776000s).
+  - `SECRET_EXPIRY` - Expiry of the keys in seconds (Default `0` = no expiry)
 
 ## Localize to your own language
 
diff --git a/storage_mem.go b/storage_mem.go
index 60239b5..2fb311c 100644
--- a/storage_mem.go
+++ b/storage_mem.go
@@ -1,20 +1,34 @@
 package main
 
-import "github.com/gofrs/uuid/v3"
+import (
+	"os"
+	"strconv"
+	"time"
+
+	"github.com/gofrs/uuid/v3"
+)
+
+type memStorageSecret struct {
+	Expiry time.Time
+	Secret string
+}
 
 type storageMem struct {
-	store map[string]string
+	store map[string]memStorageSecret
 }
 
 func newStorageMem() storage {
 	return &storageMem{
-		store: make(map[string]string),
+		store: make(map[string]memStorageSecret),
 	}
 }
 
 func (s storageMem) Create(secret string) (string, error) {
 	id := uuid.Must(uuid.NewV4()).String()
-	s.store[id] = secret
+	s.store[id] = memStorageSecret{
+		Expiry: s.expiry(),
+		Secret: secret,
+	}
 
 	return id, nil
 }
@@ -25,6 +39,25 @@ func (s storageMem) ReadAndDestroy(id string) (string, error) {
 		return "", errSecretNotFound
 	}
 
-	delete(s.store, id)
-	return secret, nil
+	defer delete(s.store, id)
+
+	if secret.Expiry.Before(time.Now()) {
+		return "", errSecretNotFound
+	}
+
+	return secret.Secret, nil
+}
+
+func (s storageMem) expiry() time.Time {
+	exp := os.Getenv("SECRET_EXPIRY")
+	if exp == "" {
+		return time.Time{}
+	}
+
+	e, err := strconv.ParseInt(exp, 10, 64)
+	if err != nil {
+		return time.Time{}
+	}
+
+	return time.Now().Add(time.Duration(e) * time.Second)
 }
diff --git a/storage_redis.go b/storage_redis.go
index 8c197e1..f48161c 100644
--- a/storage_redis.go
+++ b/storage_redis.go
@@ -65,12 +65,19 @@ func (s storageRedis) migrate() error {
 }
 
 func (s storageRedis) redisExpiry() int {
-	exp := os.Getenv("REDIS_EXPIRY")
-	if exp == "" {
+	var expStr string
+	for _, eVar := range []string{"SECRET_EXPIRY", "REDIS_EXPIRY"} {
+		if v := os.Getenv(eVar); v != "" {
+			expStr = v
+			break
+		}
+	}
+
+	if expStr == "" {
 		return 0
 	}
 
-	e, err := strconv.ParseInt(exp, 10, 64)
+	e, err := strconv.ParseInt(expStr, 10, 64)
 	if err != nil {
 		return 0
 	}