From 750de24fd1565f3c685849f1b91360cebcfa2dcc Mon Sep 17 00:00:00 2001
From: Haelnorr <lockedonohoe@gmail.com>
Date: Mon, 10 Feb 2025 22:47:20 +1100
Subject: [PATCH] Set token cookies on user login

---
 cookies/tokens.go | 40 ++++++++++++++++++++++++++++++++++------
 handlers/login.go | 14 +++++++-------
 server/routes.go  |  1 +
 3 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/cookies/tokens.go b/cookies/tokens.go
index e006756..3d2e7f5 100644
--- a/cookies/tokens.go
+++ b/cookies/tokens.go
@@ -2,13 +2,17 @@ package cookies
 
 import (
 	"net/http"
-	"projectreshoot/config"
 	"time"
+
+	"projectreshoot/config"
+	"projectreshoot/db"
+	"projectreshoot/jwt"
+
+	"github.com/pkg/errors"
 )
 
 // Get the value of the access and refresh tokens
-func GetTokens(
-	w http.ResponseWriter,
+func GetTokenStrings(
 	r *http.Request,
 ) (acc string, ref string) {
 	accCookie, accErr := r.Cookie("access")
@@ -27,22 +31,46 @@ func GetTokens(
 }
 
 // Set a token with the provided details
-func SetToken(
+func setToken(
 	w http.ResponseWriter,
-	r *http.Request,
 	config *config.Config,
 	token string,
 	scope string,
 	exp int64,
+	rememberme bool,
 ) {
 	tokenCookie := &http.Cookie{
 		Name:     scope,
 		Value:    token,
 		Path:     "/",
-		Expires:  time.Unix(exp, 0),
 		HttpOnly: true,
 		SameSite: http.SameSiteLaxMode,
 		Secure:   config.SSL,
 	}
+	if rememberme {
+		tokenCookie.Expires = time.Unix(exp, 0)
+	}
 	http.SetCookie(w, tokenCookie)
 }
+
+// Generate new tokens for the user and set them as cookies
+func SetTokenCookies(
+	w http.ResponseWriter,
+	r *http.Request,
+	config *config.Config,
+	user *db.User,
+	rememberMe bool,
+) error {
+	at, atexp, err := jwt.GenerateAccessToken(config, user, true, rememberMe)
+	if err != nil {
+		return errors.Wrap(err, "jwt.GenerateAccessToken")
+	}
+	rt, rtexp, err := jwt.GenerateRefreshToken(config, user, rememberMe)
+	if err != nil {
+		return errors.Wrap(err, "jwt.GenerateRefreshToken")
+	}
+	// Don't set the cookies until we know no errors occured
+	setToken(w, config, at, "access", atexp, rememberMe)
+	setToken(w, config, rt, "refresh", rtexp, rememberMe)
+	return nil
+}
diff --git a/handlers/login.go b/handlers/login.go
index a77b970..a7afd00 100644
--- a/handlers/login.go
+++ b/handlers/login.go
@@ -2,9 +2,9 @@ package handlers
 
 import (
 	"database/sql"
-	"fmt"
 	"net/http"
 
+	"projectreshoot/config"
 	"projectreshoot/cookies"
 	"projectreshoot/db"
 	"projectreshoot/view/component/form"
@@ -45,6 +45,7 @@ func checkRememberMe(r *http.Request) bool {
 // and on fail will return the login form again, passing the error to the
 // template for user feedback
 func HandleLoginRequest(
+	config *config.Config,
 	logger *zerolog.Logger,
 	conn *sql.DB,
 	secretKey string,
@@ -61,13 +62,12 @@ func HandleLoginRequest(
 				return
 			}
 
-			// TODO: login success, use the userID to set the session
 			rememberMe := checkRememberMe(r)
-			fmt.Printf(
-				"Login success, user: %v - remember me?: %t\n",
-				user.Username,
-				rememberMe,
-			)
+			err = cookies.SetTokenCookies(w, r, config, &user, rememberMe)
+			if err != nil {
+				form.LoginForm(err.Error()).Render(r.Context(), w)
+				logger.Warn().Caller().Err(err).Msg("Failed to set token cookies")
+			}
 
 			pageFrom := cookies.CheckPageFrom(w, r)
 			w.Header().Set("HX-Redirect", pageFrom)
diff --git a/server/routes.go b/server/routes.go
index 702bb9e..0528139 100644
--- a/server/routes.go
+++ b/server/routes.go
@@ -30,6 +30,7 @@ func addRoutes(
 	// Login page and handlers
 	mux.Handle("GET /login", handlers.HandleLoginPage(config.TrustedHost))
 	mux.Handle("POST /login", handlers.HandleLoginRequest(
+		config,
 		logger,
 		conn,
 		config.SecretKey,