From cdd98686a8b79d37874c1b9c914ae0b9696b613d Mon Sep 17 00:00:00 2001
From: Peter Sanchez <peter@netlandish.com>
Date: Thu, 29 Dec 2022 11:20:24 -0600
Subject: [PATCH] Added KeyWallet type and ability to cycle secret keys when
 signing/encrypting cookies

---
 cookies/cookies.go | 56 +++++++++++++++++++++++++++++++++-------------
 1 file changed, 41 insertions(+), 15 deletions(-)

diff --git a/cookies/cookies.go b/cookies/cookies.go
index 5e991ae..47a75ce 100644
--- a/cookies/cookies.go
+++ b/cookies/cookies.go
@@ -10,6 +10,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	mrand "math/rand"
 	"net/http"
 	"strings"
 
@@ -31,18 +32,34 @@ var (
 // key without affecting already encrypted/signed data that's
 // in the wild.
 //
-// The first key in `Keys` should be the current main key
+// # The first key in `Keys` should be the current main key
+//
+// If `Reset` is true then if any key, other than the current
+// main key was used for the cookie signing/encryption it will
+// resign/encrypt the value with the current main key. Mainly
+// useful for edge cases where you need to phase out an old
+// key for whatever reason
 type KeyWallet struct {
-	Keys [][]byte
+	Keys  [][]byte
+	Reset bool
 }
 
 // GenerateSecretKey will generate a random 32byte key used for cookie
 // signing and/or encryption
-func GenerateSecretKey() []byte {
-	key := make([]byte, 32)
+func GenerateSecretKey(keylen int, alpha bool) []byte {
+	key := make([]byte, keylen)
+	if alpha {
+		chars := `abcdefghijkmnpqrstuvwxyzABCDEFGHIJKLMNPQRSTUVWXYZ0123456789!@#$%^&*()-_=+[]{}\\|'";:/?.>,<`
+		for i := 0; i < keylen; i++ {
+			key[i] = chars[mrand.Intn(len(chars))]
+		}
+		return key
+	}
+
 	_, err := rand.Read(key)
 	if err != nil {
-		// handle error here
+		// XXX Panic?
+		panic(err)
 	}
 	return key
 }
@@ -99,8 +116,8 @@ func GetEncode(c echo.Context, name string) (*http.Cookie, error) {
 }
 
 // SetSigned will set a cookie and sign the value
-func SetSigned(c echo.Context, cookie *http.Cookie, key *KeyWallet) error {
-	mac := hmac.New(sha256.New, key.Keys[0])
+func SetSigned(c echo.Context, cookie *http.Cookie, kw *KeyWallet) error {
+	mac := hmac.New(sha256.New, kw.Keys[0])
 	mac.Write([]byte(cookie.Name))
 	mac.Write([]byte(cookie.Value))
 	signature := mac.Sum(nil)
@@ -116,7 +133,7 @@ func SetSigned(c echo.Context, cookie *http.Cookie, key *KeyWallet) error {
 }
 
 // GetSigned will get a cookie and verify it's signature
-func GetSigned(c echo.Context, name string, key *KeyWallet) (*http.Cookie, error) {
+func GetSigned(c echo.Context, name string, kw *KeyWallet) (*http.Cookie, error) {
 	cookie, err := Get(c, name)
 	if err != nil {
 		return nil, err
@@ -136,9 +153,9 @@ func GetSigned(c echo.Context, name string, key *KeyWallet) (*http.Cookie, error
 
 	// Loop through all keys
 	var found bool
-	for _, _key := range key.Keys {
+	for i, key := range kw.Keys {
 		//Recalculate the HMAC signature of the cookie name and original value.
-		mac := hmac.New(sha256.New, _key)
+		mac := hmac.New(sha256.New, key)
 		mac.Write([]byte(name))
 		mac.Write([]byte(value))
 		expectedSignature := mac.Sum(nil)
@@ -149,7 +166,12 @@ func GetSigned(c echo.Context, name string, key *KeyWallet) (*http.Cookie, error
 		if !hmac.Equal([]byte(signature), expectedSignature) {
 			continue
 		}
+
 		cookie.Value = value
+		if i != 0 && kw.Reset {
+			SetSigned(c, cookie, kw)
+			cookie.Value = value
+		}
 		found = true
 		break
 	}
@@ -162,9 +184,9 @@ func GetSigned(c echo.Context, name string, key *KeyWallet) (*http.Cookie, error
 }
 
 // SetEncrypted will set a cookie that is encrypted
-func SetEncrypted(c echo.Context, cookie *http.Cookie, key *KeyWallet) error {
+func SetEncrypted(c echo.Context, cookie *http.Cookie, kw *KeyWallet) error {
 	// Create a new AES cipher block from the secret key.
-	block, err := aes.NewCipher(key.Keys[0])
+	block, err := aes.NewCipher(kw.Keys[0])
 	if err != nil {
 		return err
 	}
@@ -205,7 +227,7 @@ func SetEncrypted(c echo.Context, cookie *http.Cookie, key *KeyWallet) error {
 }
 
 // GetEncrypted will get an encrypted cookie and decrypt it
-func GetEncrypted(c echo.Context, name string, key *KeyWallet) (*http.Cookie, error) {
+func GetEncrypted(c echo.Context, name string, kw *KeyWallet) (*http.Cookie, error) {
 	cookie, err := Get(c, name)
 	if err != nil {
 		return nil, err
@@ -216,9 +238,9 @@ func GetEncrypted(c echo.Context, name string, key *KeyWallet) (*http.Cookie, er
 	}
 
 	var found bool
-	for _, _key := range key.Keys {
+	for i, key := range kw.Keys {
 		// Create a new AES cipher block from the secret key.
-		block, err := aes.NewCipher(_key)
+		block, err := aes.NewCipher(key)
 		if err != nil {
 			// XXX Log error here?
 			continue
@@ -268,6 +290,10 @@ func GetEncrypted(c echo.Context, name string, key *KeyWallet) (*http.Cookie, er
 		}
 
 		cookie.Value = value
+		if i != 0 && kw.Reset {
+			SetEncrypted(c, cookie, kw)
+			cookie.Value = value
+		}
 		found = true
 		break
 	}
-- 
2.45.2