From 933ed6d613e04f4b7b2b448e2159c360682e7c75 Mon Sep 17 00:00:00 2001
From: Erin Shepherd <erin.shepherd@e43.eu>
Date: Wed, 8 Feb 2023 12:37:23 +0000
Subject: [PATCH] webui: add after-login redirection support

---
 internal/webui/content/login.tmpl |  1 +
 internal/webui/jwt.go             |  3 +-
 internal/webui/pg_login.go        | 50 ++++++++++++++++++++++++++++++-
 3 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/internal/webui/content/login.tmpl b/internal/webui/content/login.tmpl
index beffd2c..03e05a8 100644
--- a/internal/webui/content/login.tmpl
+++ b/internal/webui/content/login.tmpl
@@ -120,6 +120,7 @@
           </label>
         </div>
         <button class="w-100 btn btn-lg btn-primary" type="submit">Sign in</button>
+        {{if .Next}}<input type="hidden" name="next" value="{{.Next}}">{{end}}
         {{.CSRFField}}
         <!--<p class="mt-5 mb-3 text-muted">&copy; 2017–2022</p>-->
       </form>
diff --git a/internal/webui/jwt.go b/internal/webui/jwt.go
index b3ae86f..026c4c3 100644
--- a/internal/webui/jwt.go
+++ b/internal/webui/jwt.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"net/url"
 	"time"
 
 	"github.com/google/uuid"
@@ -29,7 +30,7 @@ func getUserToken(ctx context.Context) openid.Token {
 func requireLogin(w http.ResponseWriter, r *http.Request) bool {
 	tok := getUserToken(r.Context())
 	if tok == nil {
-		http.Redirect(w, r, "/login", http.StatusFound)
+		http.Redirect(w, r, "/login?next="+url.QueryEscape(r.URL.String()), http.StatusFound)
 		return false
 	}
 	return true
diff --git a/internal/webui/pg_login.go b/internal/webui/pg_login.go
index 035c4e3..3109710 100644
--- a/internal/webui/pg_login.go
+++ b/internal/webui/pg_login.go
@@ -5,6 +5,8 @@ import (
 	"errors"
 	"html/template"
 	"net/http"
+	"net/url"
+	"strings"
 
 	"github.com/gorilla/csrf"
 	"go.e43.eu/authentricity/internal/models"
@@ -36,12 +38,19 @@ func (fr loginFailureReason) Error() string {
 }
 
 func (s *Service) loginGet(w http.ResponseWriter, r *http.Request) {
+	L := zap.L()
 	tok := getUserToken(r.Context())
 	if tok != nil {
 		http.Redirect(w, r, "/", http.StatusFound)
 		return
 	}
 
+	if err := r.ParseForm(); err != nil {
+		L.Error("Error parsing form data", zap.Error(err))
+		s.renderError(w)
+		return
+	}
+
 	s.showLoginPage(w, r, "")
 }
 
@@ -83,7 +92,44 @@ func (s *Service) loginPost(w http.ResponseWriter, r *http.Request) {
 
 	cookie := s.buildTokenCookie(serialized, 86400)
 	http.SetCookie(w, &cookie)
-	http.Redirect(w, r, "/entity/"+user.UUID.String(), http.StatusFound)
+
+	nextURL, ok := s.getLoginNextURL(r)
+	if !ok {
+		nextURL = "/entity/" + user.UUID.String()
+	}
+
+	http.Redirect(w, r, nextURL, http.StatusFound)
+}
+
+func (s *Service) getLoginNextURL(r *http.Request) (string, bool) {
+	L := zap.L()
+
+	next := r.PostForm.Get("next")
+	if next == "" {
+		L.Debug("Ignoring next as empty or unset")
+		return "", false
+	}
+
+	nextURL, err := url.ParseRequestURI(next)
+	switch {
+	case err != nil:
+		L.Debug("Ignoring next as failed to parse",
+			zap.String("next", next), zap.Error(err))
+		return "", false
+
+	case nextURL.Scheme != "https":
+		L.Debug("Ignoring next as scheme not HTTPS",
+			zap.String("next", next))
+	case nextURL.Host != "" &&
+		nextURL.Host != s.cookieDomain &&
+		!strings.HasSuffix(nextURL.Host, "."+s.cookieDomain):
+		L.Debug("Ignoring next as not within cookie domain",
+			zap.String("next", next))
+		return "", false
+	}
+
+	// We re-stringify the URL as this reduces request smuggling type risks
+	return nextURL.String(), true
 }
 
 func (s *Service) tryLogin(ctx context.Context, username, password string) (*models.UserRecord, error) {
@@ -134,10 +180,12 @@ func (s *Service) showLoginPage(w http.ResponseWriter, r *http.Request, message
 		ShowError    bool
 		ErrorMessage string
 		CSRFField    template.HTML
+		Next         string
 	}{
 		ShowError:    message != "",
 		ErrorMessage: message,
 		CSRFField:    csrf.TemplateField(r),
+		Next:         r.Form.Get("next"),
 	}
 
 	err := s.templates.ExecuteTemplate(w, "login.tmpl", params)