dkforest

commit b4b5b6fbc1c76f1530ad5dd07402c13ed6dfa579
parent d196f046d034f3e26191882fdf4f4fa8cef8f3a2
Author: n0tr1v <n0tr1v@protonmail.com>
Date:   Sun, 18 Dec 2022 01:36:27 -0800

simplify code by having utils functions to handle encrypted cookies

Diffstat:
M
pkg/web/handlers/handlers.go
 | 
30
++++++------------------------
M
pkg/web/handlers/utils/utils.go
 | 
38
++++++++++++++++++++++++++++++++++++++

2 files changed, 44 insertions(+), 24 deletions(-)
diff --git a/pkg/web/handlers/handlers.go b/pkg/web/handlers/handlers.go
@@ -768,7 +768,8 @@ func tmpSignupHandler(c echo.Context) error {
 func waitPageWrapper(c echo.Context, clb echo.HandlerFunc, cookieName string) error {
 	start := time.Now().UnixNano()
 	var signupToken string
-	if cc, err := c.Cookie(cookieName); err != nil {
+
+	if payload, err := hutils.EncCookie[map[string]string](c, cookieName); err != nil {
 		// No cookie found, we create one and display the waiting page.
 		waitTime := utils.Random(5, 15)
 		signupToken = utils.GenerateToken10()
@@ -778,10 +779,7 @@ func waitPageWrapper(c echo.Context, clb echo.HandlerFunc, cookieName string) er
 			"now":   utils.FormatInt64(time.Now().UnixMilli()),
 			"unix":  utils.FormatInt64(time.Now().Unix() + waitTime - 1), // unix time at which the wait time is over
 		}
-		by, _ := json.Marshal(payload)
-		encryptedVal, _ := utils.EncryptAES(by, []byte(config.Global.MasterKey()))
-		valB64 := base64.URLEncoding.EncodeToString(encryptedVal)
-		c.SetCookie(hutils.CreateCookie(cookieName, valB64, utils.OneMinuteSecs*5))
+		c.SetCookie(hutils.CreateEncCookie(cookieName, payload, utils.OneMinuteSecs*5))
 
 		var data1 waitData
 		// Generate css frames
@@ -792,19 +790,6 @@ func waitPageWrapper(c echo.Context, clb echo.HandlerFunc, cookieName string) er
 
 	} else {
 		// Cookie was found, incr counter then call callback
-		valB64 := cc.Value
-		val, err := base64.URLEncoding.DecodeString(valB64)
-		if err != nil {
-			return c.String(http.StatusFound, "Invalid token 1")
-		}
-		v, err := utils.DecryptAES(val, []byte(config.Global.MasterKey()))
-		if err != nil {
-			return c.String(http.StatusFound, "Invalid token 3")
-		}
-		var payload map[string]string
-		if err := json.Unmarshal(v, &payload); err != nil {
-			return c.String(http.StatusFound, "Invalid token 2")
-		}
 		signupToken = payload["token"]
 		start = utils.DoParseInt64(payload["now"])
 		if c.Request().Method == http.MethodGet {
@@ -822,18 +807,15 @@ func waitPageWrapper(c echo.Context, clb echo.HandlerFunc, cookieName string) er
 
 			// If the wait time is over, and you reload the protected page more than 4 times, we make you wait 1min
 			if count >= 4 {
-				c.SetCookie(hutils.CreateCookie(cookieName, valB64, utils.OneMinuteSecs))
-				return c.String(http.StatusFound, "You tried to reload the page to many times. Now you have to wait one minute.")
+				c.SetCookie(hutils.CreateEncCookie(cookieName, payload, utils.OneMinuteSecs))
+				return c.String(http.StatusFound, "You tried to reload the page too many times. Now you have to wait one minute.")
 			}
 			newPayload := map[string]string{
 				"count": utils.FormatInt64(count + 1),
 				"now":   utils.FormatInt64(time.Now().UnixMilli()),
 				"token": signupToken,
 			}
-			by, _ := json.Marshal(newPayload)
-			newEncryptedVal, _ := utils.EncryptAES(by, []byte(config.Global.MasterKey()))
-			newValB64 := base64.URLEncoding.EncodeToString(newEncryptedVal)
-			c.SetCookie(hutils.CreateCookie(cookieName, newValB64, utils.OneMinuteSecs*5))
+			c.SetCookie(hutils.CreateEncCookie(cookieName, newPayload, utils.OneMinuteSecs*5))
 		}
 	}
 	c.Set("start", start)
diff --git a/pkg/web/handlers/utils/utils.go b/pkg/web/handlers/utils/utils.go
@@ -2,6 +2,8 @@ package utils
 
 import (
 	"dkforest/pkg/captcha"
+	"encoding/base64"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"net/http"
@@ -38,6 +40,42 @@ func CreateCookie(name, value string, maxAge int64) *http.Cookie {
 	return cookie
 }
 
+// CreateEncCookie return a cookie where the value has been json marshaled, encrypted and base64 encoded
+func CreateEncCookie(name string, value any, maxAge int64) *http.Cookie {
+	by, err := json.Marshal(value)
+	if err != nil {
+		return nil
+	}
+	encryptedVal, err := utils.EncryptAES(by, []byte(config.Global.MasterKey()))
+	if err != nil {
+		return nil
+	}
+	valB64 := base64.URLEncoding.EncodeToString(encryptedVal)
+	return CreateCookie(name, valB64, maxAge)
+}
+
+// EncCookie gets back the value of an encrypted cookie
+func EncCookie[T any](c echo.Context, name string) (T, error) {
+	var zero T
+	cc, err := c.Cookie(name)
+	if err != nil {
+		return zero, err
+	}
+	val, err := base64.URLEncoding.DecodeString(cc.Value)
+	if err != nil {
+		return zero, err
+	}
+	v, err := utils.DecryptAES(val, []byte(config.Global.MasterKey()))
+	if err != nil {
+		return zero, err
+	}
+	var out T
+	if err := json.Unmarshal(v, &out); err != nil {
+		return zero, err
+	}
+	return out, nil
+}
+
 func DeleteCookie(name string) *http.Cookie {
 	return CreateCookie(name, "", -1)
 }