Add redirect back to specific page

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

pkg/saferedirect/saferedirect.go

@@ -0,0 +1,73 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package saferedirect
+
+import (
+	"net/http"
+	"net/url"
+	"strings"
+)
+
+type (
+	SafeRedirect struct {
+		AllowedHost string
+	}
+)
+
+func (sr *SafeRedirect) Validate(redirectURL string) (string, bool) {
+	if redirectURL == "" {
+		return "", false
+	}
+
+	if strings.HasPrefix(redirectURL, "/") {
+		return redirectURL, true
+	}
+
+	parsedURL, err := url.Parse(redirectURL)
+	if err != nil {
+		return "", false
+	}
+
+	if parsedURL.IsAbs() {
+		if parsedURL.Scheme != "http" && parsedURL.Scheme != "https" {
+			return "", false
+		}
+
+		if sr.AllowedHost != "" && parsedURL.Host != sr.AllowedHost {
+			return "", false
+		}
+
+		return redirectURL, true
+	}
+
+	return "", false
+}
+
+func (sr *SafeRedirect) GetSafeRedirectURL(redirectURL, fallbackURL string) string {
+	if safeURL, isValid := sr.Validate(redirectURL); isValid {
+		return safeURL
+	}
+	return fallbackURL
+}
+
+func (sr *SafeRedirect) Redirect(w http.ResponseWriter, r *http.Request, redirectURL, fallbackURL string, statusCode int) {
+	safeURL := sr.GetSafeRedirectURL(redirectURL, fallbackURL)
+	http.Redirect(w, r, safeURL, statusCode)
+}
+
+func (sr *SafeRedirect) RedirectFromQuery(w http.ResponseWriter, r *http.Request, paramName, fallbackURL string, statusCode int) {
+	redirectURL := r.URL.Query().Get(paramName)
+	sr.Redirect(w, r, redirectURL, fallbackURL, statusCode)
+}