login: remember redirects

oauth.go

@@ -7,6 +7,7 @@ import (
 	"encoding/json"
 	"html/template"
 	"net/http"
+	"net/url"
 	"sync"
 	"time"
 
@@ -81,8 +82,8 @@ func (s *OAuthStore) DeleteSession(sessionID string) {
 	s.mutex.Unlock()
 }
 
-func sendToLoginPage(w http.ResponseWriter, r *http.Request) {
-	http.Redirect(w, r, "/oauth/login", http.StatusTemporaryRedirect)
+func sendToLoginPage(w http.ResponseWriter, r *http.Request, origin string) {
+	http.Redirect(w, r, "/oauth/login?source="+url.QueryEscape(origin), http.StatusTemporaryRedirect)
 }
 
 func sendToUnauthorized(w http.ResponseWriter, r *http.Request) {
@@ -95,23 +96,12 @@ 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")
+func getRedirectFromCookie(r *http.Request) (string, error) {
+	cookie, err := r.Cookie("redirect_origin")
 	if err != nil {
-		return "/"
+		return "", err
 	}
-	return cookie.Value
+	return cookie.Value, nil
 }
 
 //go:embed templates/LoginPage.html
@@ -147,6 +137,17 @@ func (s *OAuthStore) LoginPage() http.Handler {
 			})
 		}
 
+		if r.URL.Query().Has("source") {
+			http.SetCookie(w, &http.Cookie{
+				Name:     "redirect_origin",
+				Value:    r.URL.Query().Get("source"),
+				HttpOnly: true,
+				Secure:   true,
+				SameSite: http.SameSiteLaxMode,
+				MaxAge:   60 * 10,
+			})
+		}
+
 		url := s.oa2.AuthCodeURL(state)
 
 		provider := s.config.OAuthProvider.Kind
@@ -173,13 +174,12 @@ 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)
+			sendToLoginPage(w, r, r.URL.Path)
 			return
 		}
 		sess, exists := s.GetSession(cookie.Value)
 		if !exists {
-			setRedirectCookie(w, r.URL.Path)
-			sendToLoginPage(w, r)
+			sendToLoginPage(w, r, r.URL.Path)
 			return
 		}
 
@@ -232,7 +232,10 @@ func (s *OAuthStore) CallbackHandler() http.Handler {
 			Path:     "/",
 		})
 
-		redirect := getRedirectFromCookie(r)
+		redirect, err := getRedirectFromCookie(r)
+		if err != nil {
+			redirect = "/"
+		}
 
 		// clear cookies
 		http.SetCookie(w, &http.Cookie{
@@ -246,7 +249,6 @@ func (s *OAuthStore) CallbackHandler() http.Handler {
 			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, redirect, http.StatusTemporaryRedirect)
 	})
 }