Add redirect back to specific page

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

pkg/probod/probod.go

@@ -31,6 +31,7 @@ import (
 	"github.com/getprobo/probo/pkg/crypto/passwdhash"
 	"github.com/getprobo/probo/pkg/mailer"
 	"github.com/getprobo/probo/pkg/probo"
+	"github.com/getprobo/probo/pkg/saferedirect"
 	"github.com/getprobo/probo/pkg/server"
 	console_v1 "github.com/getprobo/probo/pkg/server/api/console/v1"
 	"github.com/getprobo/probo/pkg/usrmgr"
@@ -207,6 +208,7 @@ func (impl *Implm) Run(
 			Probo:             proboService,
 			Usrmgr:            usrmgrService,
 			ConnectorRegistry: defaultConnectorRegistry,
+			SafeRedirect:      &saferedirect.SafeRedirect{AllowedHost: impl.cfg.Hostname},
 			Auth: console_v1.AuthConfig{
 				CookieName:      impl.cfg.Auth.Cookie.Name,
 				CookieDomain:    impl.cfg.Auth.Cookie.Domain,

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)
+}

pkg/saferedirect/saferedirect_test.go

@@ -0,0 +1,292 @@
+// 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_test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/getprobo/probo/pkg/saferedirect"
+)
+
+func TestSafeRedirect_Validate(t *testing.T) {
+	tests := []struct {
+		name            string
+		allowedHost     string
+		redirectURL     string
+		expectedURL     string
+		expectedIsValid bool
+	}{
+		{
+			name:            "empty redirect URL",
+			allowedHost:     "example.com",
+			redirectURL:     "",
+			expectedURL:     "",
+			expectedIsValid: false,
+		},
+		{
+			name:            "relative URL",
+			allowedHost:     "example.com",
+			redirectURL:     "/dashboard",
+			expectedURL:     "/dashboard",
+			expectedIsValid: true,
+		},
+		{
+			name:            "allowed absolute URL",
+			allowedHost:     "example.com",
+			redirectURL:     "https://example.com/dashboard",
+			expectedURL:     "https://example.com/dashboard",
+			expectedIsValid: true,
+		},
+		{
+			name:            "disallowed host",
+			allowedHost:     "example.com",
+			redirectURL:     "https://evil.com/phishing",
+			expectedURL:     "",
+			expectedIsValid: false,
+		},
+		{
+			name:            "disallowed scheme (javascript:)",
+			allowedHost:     "example.com",
+			redirectURL:     "javascript:alert('xss')",
+			expectedURL:     "",
+			expectedIsValid: false,
+		},
+		{
+			name:            "disallowed scheme (data:)",
+			allowedHost:     "example.com",
+			redirectURL:     "data:text/html;base64,PHNjcmlwdD5hbGVydCgnWFNTJyk8L3NjcmlwdD4=",
+			expectedURL:     "",
+			expectedIsValid: false,
+		},
+		{
+			name:            "no allowed host restriction",
+			allowedHost:     "",
+			redirectURL:     "https://any-domain.com/page",
+			expectedURL:     "https://any-domain.com/page",
+			expectedIsValid: true,
+		},
+		{
+			name:            "invalid URL",
+			allowedHost:     "example.com",
+			redirectURL:     "https://[invalid-url",
+			expectedURL:     "",
+			expectedIsValid: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			sr := saferedirect.SafeRedirect{
+				AllowedHost: tt.allowedHost,
+			}
+
+			gotURL, gotIsValid := sr.Validate(tt.redirectURL)
+			if gotIsValid != tt.expectedIsValid {
+				t.Errorf("Validate() isValid = %v, want %v", gotIsValid, tt.expectedIsValid)
+			}
+			if gotURL != tt.expectedURL {
+				t.Errorf("Validate() url = %v, want %v", gotURL, tt.expectedURL)
+			}
+		})
+	}
+}
+
+func TestSafeRedirect_GetSafeRedirectURL(t *testing.T) {
+	tests := []struct {
+		name        string
+		allowedHost string
+		redirectURL string
+		fallbackURL string
+		expectedURL string
+	}{
+		{
+			name:        "safe redirect URL",
+			allowedHost: "example.com",
+			redirectURL: "/dashboard",
+			fallbackURL: "/home",
+			expectedURL: "/dashboard",
+		},
+		{
+			name:        "unsafe redirect URL",
+			allowedHost: "example.com",
+			redirectURL: "https://evil.com/phishing",
+			fallbackURL: "/home",
+			expectedURL: "/home",
+		},
+		{
+			name:        "empty redirect URL",
+			allowedHost: "example.com",
+			redirectURL: "",
+			fallbackURL: "/home",
+			expectedURL: "/home",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			sr := saferedirect.SafeRedirect{
+				AllowedHost: tt.allowedHost,
+			}
+
+			gotURL := sr.GetSafeRedirectURL(tt.redirectURL, tt.fallbackURL)
+			if gotURL != tt.expectedURL {
+				t.Errorf("GetSafeRedirectURL() = %v, want %v", gotURL, tt.expectedURL)
+			}
+		})
+	}
+}
+
+func TestSafeRedirect_Redirect(t *testing.T) {
+	tests := []struct {
+		name           string
+		allowedHost    string
+		redirectURL    string
+		fallbackURL    string
+		expectedStatus int
+		expectedURL    string
+	}{
+		{
+			name:           "safe redirect URL",
+			allowedHost:    "example.com",
+			redirectURL:    "/dashboard",
+			fallbackURL:    "/home",
+			expectedStatus: http.StatusFound,
+			expectedURL:    "/dashboard",
+		},
+		{
+			name:           "unsafe redirect URL",
+			allowedHost:    "example.com",
+			redirectURL:    "https://evil.com/phishing",
+			fallbackURL:    "/home",
+			expectedStatus: http.StatusFound,
+			expectedURL:    "/home",
+		},
+		{
+			name:           "empty redirect URL",
+			allowedHost:    "example.com",
+			redirectURL:    "",
+			fallbackURL:    "/home",
+			expectedStatus: http.StatusFound,
+			expectedURL:    "/home",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			sr := saferedirect.SafeRedirect{
+				AllowedHost: tt.allowedHost,
+			}
+
+			// Create a test HTTP recorder to capture the response
+			w := httptest.NewRecorder()
+			r := httptest.NewRequest("GET", "http://test.com", nil)
+
+			sr.Redirect(w, r, tt.redirectURL, tt.fallbackURL, tt.expectedStatus)
+
+			// Check that we got the expected status code
+			if w.Code != tt.expectedStatus {
+				t.Errorf("Redirect() status = %v, want %v", w.Code, tt.expectedStatus)
+			}
+
+			// Check that the Location header contains the expected URL
+			location := w.Header().Get("Location")
+			if location != tt.expectedURL {
+				t.Errorf("Redirect() location = %v, want %v", location, tt.expectedURL)
+			}
+		})
+	}
+}
+
+func TestSafeRedirect_RedirectFromQuery(t *testing.T) {
+	tests := []struct {
+		name           string
+		allowedHost    string
+		queryParam     string
+		queryValue     string
+		fallbackURL    string
+		expectedStatus int
+		expectedURL    string
+	}{
+		{
+			name:           "safe continue param",
+			allowedHost:    "example.com",
+			queryParam:     "continue",
+			queryValue:     "/dashboard",
+			fallbackURL:    "/home",
+			expectedStatus: http.StatusFound,
+			expectedURL:    "/dashboard",
+		},
+		{
+			name:           "unsafe continue param",
+			allowedHost:    "example.com",
+			queryParam:     "continue",
+			queryValue:     "https://evil.com/phishing",
+			fallbackURL:    "/home",
+			expectedStatus: http.StatusFound,
+			expectedURL:    "/home",
+		},
+		{
+			name:           "missing continue param",
+			allowedHost:    "example.com",
+			queryParam:     "continue",
+			queryValue:     "",
+			fallbackURL:    "/home",
+			expectedStatus: http.StatusFound,
+			expectedURL:    "/home",
+		},
+		{
+			name:           "different query param name",
+			allowedHost:    "example.com",
+			queryParam:     "next",
+			queryValue:     "/profile",
+			fallbackURL:    "/home",
+			expectedStatus: http.StatusFound,
+			expectedURL:    "/profile",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			sr := saferedirect.SafeRedirect{
+				AllowedHost: tt.allowedHost,
+			}
+
+			// Create a test HTTP recorder to capture the response
+			w := httptest.NewRecorder()
+
+			// Create request with query parameter
+			url := "http://test.com"
+			if tt.queryValue != "" {
+				url = "http://test.com?" + tt.queryParam + "=" + tt.queryValue
+			}
+			r := httptest.NewRequest("GET", url, nil)
+
+			sr.RedirectFromQuery(w, r, tt.queryParam, tt.fallbackURL, tt.expectedStatus)
+
+			// Check that we got the expected status code
+			if w.Code != tt.expectedStatus {
+				t.Errorf("RedirectFromQuery() status = %v, want %v", w.Code, tt.expectedStatus)
+			}
+
+			// Check that the Location header contains the expected URL
+			location := w.Header().Get("Location")
+			if location != tt.expectedURL {
+				t.Errorf("RedirectFromQuery() location = %v, want %v", location, tt.expectedURL)
+			}
+		})
+	}
+}

pkg/server/api/api.go

@@ -20,6 +20,7 @@ import (
 
 	"github.com/getprobo/probo/pkg/connector"
 	"github.com/getprobo/probo/pkg/probo"
+	"github.com/getprobo/probo/pkg/saferedirect"
 	console_v1 "github.com/getprobo/probo/pkg/server/api/console/v1"
 	"github.com/getprobo/probo/pkg/usrmgr"
 	"github.com/go-chi/chi/v5"
@@ -34,6 +35,7 @@ type (
 		Usrmgr            *usrmgr.Service
 		Auth              console_v1.AuthConfig
 		ConnectorRegistry *connector.ConnectorRegistry
+		SafeRedirect      *saferedirect.SafeRedirect
 	}
 
 	Server struct {
@@ -103,7 +105,7 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	router.Use(cors.Handler(corsOpts))
 
 	// Mount the console API with authentication
-	router.Mount("/console/v1", console_v1.NewMux(s.cfg.Probo, s.cfg.Usrmgr, s.cfg.Auth, s.cfg.ConnectorRegistry))
+	router.Mount("/console/v1", console_v1.NewMux(s.cfg.Probo, s.cfg.Usrmgr, s.cfg.Auth, s.cfg.ConnectorRegistry, s.cfg.SafeRedirect))
 
 	router.ServeHTTP(w, r)
 }

pkg/server/api/console/v1/resolver.go

@@ -32,6 +32,7 @@ import (
 	"github.com/getprobo/probo/pkg/coredata"
 	"github.com/getprobo/probo/pkg/gid"
 	"github.com/getprobo/probo/pkg/probo"
+	"github.com/getprobo/probo/pkg/saferedirect"
 	"github.com/getprobo/probo/pkg/securecookie"
 	"github.com/getprobo/probo/pkg/server/api/console/v1/schema"
 	"github.com/getprobo/probo/pkg/usrmgr"
@@ -73,7 +74,7 @@ func UserFromContext(ctx context.Context) *coredata.User {
 	return user
 }
 
-func NewMux(proboSvc *probo.Service, usrmgrSvc *usrmgr.Service, authCfg AuthConfig, connectorRegistry *connector.ConnectorRegistry) *chi.Mux {
+func NewMux(proboSvc *probo.Service, usrmgrSvc *usrmgr.Service, authCfg AuthConfig, connectorRegistry *connector.ConnectorRegistry, safeRedirect *saferedirect.SafeRedirect) *chi.Mux {
 	r := chi.NewMux()
 
 	r.Post("/auth/register", SignUpHandler(usrmgrSvc, authCfg))
@@ -137,7 +138,7 @@ func NewMux(proboSvc *probo.Service, usrmgrSvc *usrmgr.Service, authCfg AuthConf
 			panic(fmt.Errorf("failed to create or update connector: %w", err))
 		}
 
-		http.Redirect(w, r, "/foo", http.StatusSeeOther)
+		safeRedirect.RedirectFromQuery(w, r, "continue", "/", http.StatusSeeOther)
 	}))
 
 	r.Get("/", playground.Handler("GraphQL", "/api/console/v1/query"))

pkg/server/server.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/getprobo/probo/pkg/connector"
 	"github.com/getprobo/probo/pkg/probo"
+	"github.com/getprobo/probo/pkg/saferedirect"
 	"github.com/getprobo/probo/pkg/server/api"
 	console_v1 "github.com/getprobo/probo/pkg/server/api/console/v1"
 	"github.com/getprobo/probo/pkg/server/web"
@@ -35,6 +36,7 @@ type Config struct {
 	Usrmgr            *usrmgr.Service
 	Auth              console_v1.AuthConfig
 	ConnectorRegistry *connector.ConnectorRegistry
+	SafeRedirect      *saferedirect.SafeRedirect
 }
 
 // Server represents the main server that handles both API and frontend requests
@@ -53,6 +55,7 @@ func NewServer(cfg Config) (*Server, error) {
 		Usrmgr:            cfg.Usrmgr,
 		Auth:              cfg.Auth,
 		ConnectorRegistry: cfg.ConnectorRegistry,
+		SafeRedirect:      cfg.SafeRedirect,
 	}
 	apiServer, err := api.NewServer(apiCfg)
 	if err != nil {