From a617298b5d0ec445e79387bfa408828c3cddf6c7 Mon Sep 17 00:00:00 2001
From: Lukas Werner <lukas.m.werner@outlook.com>
Date: Sun, 31 Aug 2025 10:20:03 -0700
Subject: [PATCH] feat: add proper redirects

---
 oauth.go | 34 +++++++++++++++++++++++++++++++---
 1 file changed, 31 insertions(+), 3 deletions(-)

diff --git a/oauth.go b/oauth.go
index d622102..974e70f 100644
--- a/oauth.go
+++ b/oauth.go
@@ -7,6 +7,8 @@ import (
 	"encoding/json"
 	"html/template"
 	"net/http"
+	"net/http/cookiejar"
+	"net/url"
 	"sync"
 	"time"
 
@@ -95,6 +97,25 @@ func generateRandomToken() string {
 	return base64.StdEncoding.EncodeToString(b)
 }
 
+func setRedirectCookie(w http.ResponseWriter, path string) {
+	http.SetCookie(w,
+		&http.Cookie{
+			Name:     "redirect_on_completion",
+			Value:    path,
+			Path:     "/",
+			HttpOnly: true,
+			SameSite: http.SameSiteStrictMode,
+		})
+}
+
+func getRedirectFromCookie(r *http.Request) string {
+	cookie, err := r.Cookie("redirect_on_completion")
+	if err != nil {
+		return "/"
+	}
+	return cookie.Value
+}
+
 //go:embed templates/LoginPage.html
 var loginPageContent string
 
@@ -152,7 +173,6 @@ func (s *OAuthStore) LoginPage() http.Handler {
 
 func (s *OAuthStore) Protected(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-
 		cookie, err := r.Cookie(SessionCookie)
 		if err != nil {
 			sendToLoginPage(w, r)
@@ -160,6 +180,7 @@ func (s *OAuthStore) Protected(next http.Handler) http.Handler {
 		}
 		sess, exists := s.GetSession(cookie.Value)
 		if !exists {
+			setRedirectCookie(w, r.URL.Path)
 			sendToLoginPage(w, r)
 			return
 		}
@@ -213,15 +234,22 @@ func (s *OAuthStore) CallbackHandler() http.Handler {
 			Path:     "/",
 		})
 
-		// clear cookie
+		redirect := getRedirectFromCookie(r)
+
+		// clear cookies
 		http.SetCookie(w, &http.Cookie{
 			Name:   "oauth_state",
 			Value:  "",
 			MaxAge: -1,
 		})
+		http.SetCookie(w, &http.Cookie{
+			Name:   "redirect_on_completion",
+			Value:  "",
+			MaxAge: -1,
+		})
 
 		// TODO: remember what path the user was on and redirect them back there after doing the whole login process
-		http.Redirect(w, r, "/", http.StatusTemporaryRedirect)
+		http.Redirect(w, r, redirect, http.StatusTemporaryRedirect)
 	})
 }