Fix compliance page login redirect to custom domains

SafeRedirect previously matched against a single static host string,
so OIDC callbacks always fell back to the console instead of
redirecting back to compliance pages on custom domains. Refactor
AllowedHost into a dynamic AllowedHostFunc and wire a trust-service
lookup into the connect handler so custom domain hosts are accepted.

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

pkg/saferedirect/saferedirect.go

@@ -15,18 +15,40 @@
 package saferedirect
 
 import (
+	"context"
 	"net/http"
 	"net/url"
 	"strings"
 )
 
 type (
+	AllowedHostFunc func(ctx context.Context, host string) bool
+
 	SafeRedirect struct {
-		AllowedHost string
+		allowedHost AllowedHostFunc
 	}
 )
 
-func (sr *SafeRedirect) Validate(redirectURL string) (string, bool) {
+func New(allowedHost AllowedHostFunc) *SafeRedirect {
+	return &SafeRedirect{allowedHost: allowedHost}
+}
+
+// StaticHosts returns an AllowedHost function that matches against a fixed
+// list of hosts.
+func StaticHosts(hosts ...string) AllowedHostFunc {
+	allowed := make(map[string]bool, len(hosts))
+	for _, h := range hosts {
+		if h != "" {
+			allowed[h] = true
+		}
+	}
+
+	return func(_ context.Context, host string) bool {
+		return allowed[host]
+	}
+}
+
+func (sr *SafeRedirect) Validate(ctx context.Context, redirectURL string) (string, bool) {
 	if redirectURL == "" {
 		return "", false
 	}
@@ -48,7 +70,7 @@ func (sr *SafeRedirect) Validate(redirectURL string) (string, bool) {
 			return "", false
 		}
 
-		if sr.AllowedHost != "" && parsedURL.Host != sr.AllowedHost {
+		if sr.allowedHost != nil && !sr.allowedHost(ctx, parsedURL.Host) {
 			return "", false
 		}
 
@@ -58,14 +80,14 @@ func (sr *SafeRedirect) Validate(redirectURL string) (string, bool) {
 	return "", false
 }
 
-func (sr *SafeRedirect) GetSafeRedirectURL(redirectURL, fallbackURL string) string {
+func (sr *SafeRedirect) GetSafeRedirectURL(ctx context.Context, redirectURL, fallbackURL string) string {
-	if safeURL, isValid := sr.Validate(redirectURL); isValid {
+	if safeURL, isValid := sr.Validate(ctx, 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)
+	safeURL := sr.GetSafeRedirectURL(r.Context(), redirectURL, fallbackURL)
 	http.Redirect(w, r, safeURL, statusCode)
 }

pkg/saferedirect/saferedirect_test.go

@@ -15,6 +15,7 @@
 package saferedirect_test
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -23,79 +24,81 @@ import (
 )
 
 func TestSafeRedirect_Validate(t *testing.T) {
+	t.Parallel()
+
 	tests := []struct {
 		name            string
-		allowedHost     string
+		allowedHost     saferedirect.AllowedHostFunc
 		redirectURL     string
 		expectedURL     string
 		expectedIsValid bool
 	}{
 		{
 			name:            "empty redirect URL",
-			allowedHost:     "example.com",
+			allowedHost:     saferedirect.StaticHosts("example.com"),
 			redirectURL:     "",
 			expectedURL:     "",
 			expectedIsValid: false,
 		},
 		{
 			name:            "relative URL",
-			allowedHost:     "example.com",
+			allowedHost:     saferedirect.StaticHosts("example.com"),
 			redirectURL:     "/dashboard",
 			expectedURL:     "/dashboard",
 			expectedIsValid: true,
 		},
 		{
 			name:            "allowed absolute URL",
-			allowedHost:     "example.com",
+			allowedHost:     saferedirect.StaticHosts("example.com"),
 			redirectURL:     "https://example.com/dashboard",
 			expectedURL:     "https://example.com/dashboard",
 			expectedIsValid: true,
 		},
 		{
 			name:            "disallowed host",
-			allowedHost:     "example.com",
+			allowedHost:     saferedirect.StaticHosts("example.com"),
 			redirectURL:     "https://evil.com/phishing",
 			expectedURL:     "",
 			expectedIsValid: false,
 		},
 		{
 			name:            "disallowed scheme (javascript:)",
-			allowedHost:     "example.com",
+			allowedHost:     saferedirect.StaticHosts("example.com"),
 			redirectURL:     "javascript:alert('xss')",
 			expectedURL:     "",
 			expectedIsValid: false,
 		},
 		{
 			name:            "disallowed scheme (data:)",
-			allowedHost:     "example.com",
+			allowedHost:     saferedirect.StaticHosts("example.com"),
 			redirectURL:     "data:text/html;base64,PHNjcmlwdD5hbGVydCgnWFNTJyk8L3NjcmlwdD4=",
 			expectedURL:     "",
 			expectedIsValid: false,
 		},
 		{
 			name:            "no allowed host restriction",
-			allowedHost:     "",
+			allowedHost:     nil,
 			redirectURL:     "https://any-domain.com/page",
 			expectedURL:     "https://any-domain.com/page",
 			expectedIsValid: true,
 		},
 		{
 			name:            "invalid URL",
-			allowedHost:     "example.com",
+			allowedHost:     saferedirect.StaticHosts("example.com"),
 			redirectURL:     "https://[invalid-url",
 			expectedURL:     "",
 			expectedIsValid: false,
 		},
 		{
 			name:            "double slash attack",
-			allowedHost:     "example.com",
+			allowedHost:     saferedirect.StaticHosts("example.com"),
 			redirectURL:     "//evil.com/phishing",
 			expectedURL:     "",
 			expectedIsValid: false,
 		},
 		{
 			name:            "slash-backslash attack",
-			allowedHost:     "example.com",
+			allowedHost:     saferedirect.StaticHosts("example.com"),
 			redirectURL:     "/\\evil.com/phishing",
 			expectedURL:     "",
 			expectedIsValid: false,
@@ -104,11 +107,11 @@ func TestSafeRedirect_Validate(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			sr := saferedirect.SafeRedirect{
+			t.Parallel()
-				AllowedHost: tt.allowedHost,
-			}
 
-			gotURL, gotIsValid := sr.Validate(tt.redirectURL)
+			sr := saferedirect.New(tt.allowedHost)
+
+			gotURL, gotIsValid := sr.Validate(context.Background(), tt.redirectURL)
 			if gotIsValid != tt.expectedIsValid {
 				t.Errorf("Validate() isValid = %v, want %v", gotIsValid, tt.expectedIsValid)
 			}
@@ -120,44 +123,46 @@ func TestSafeRedirect_Validate(t *testing.T) {
 }
 
 func TestSafeRedirect_GetSafeRedirectURL(t *testing.T) {
+	t.Parallel()
+
 	tests := []struct {
 		name        string
-		allowedHost string
+		allowedHost saferedirect.AllowedHostFunc
 		redirectURL string
 		fallbackURL string
 		expectedURL string
 	}{
 		{
 			name:        "safe redirect URL",
-			allowedHost: "example.com",
+			allowedHost: saferedirect.StaticHosts("example.com"),
 			redirectURL: "/dashboard",
 			fallbackURL: "/home",
 			expectedURL: "/dashboard",
 		},
 		{
 			name:        "unsafe redirect URL",
-			allowedHost: "example.com",
+			allowedHost: saferedirect.StaticHosts("example.com"),
 			redirectURL: "https://evil.com/phishing",
 			fallbackURL: "/home",
 			expectedURL: "/home",
 		},
 		{
 			name:        "empty redirect URL",
-			allowedHost: "example.com",
+			allowedHost: saferedirect.StaticHosts("example.com"),
 			redirectURL: "",
 			fallbackURL: "/home",
 			expectedURL: "/home",
 		},
 		{
 			name:        "double slash attack",
-			allowedHost: "example.com",
+			allowedHost: saferedirect.StaticHosts("example.com"),
 			redirectURL: "//evil.com/phishing",
 			fallbackURL: "/home",
 			expectedURL: "/home",
 		},
 		{
 			name:        "slash-backslash attack",
-			allowedHost: "example.com",
+			allowedHost: saferedirect.StaticHosts("example.com"),
 			redirectURL: "/\\evil.com/phishing",
 			fallbackURL: "/home",
 			expectedURL: "/home",
@@ -166,11 +171,11 @@ func TestSafeRedirect_GetSafeRedirectURL(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			sr := saferedirect.SafeRedirect{
+			t.Parallel()
-				AllowedHost: tt.allowedHost,
-			}
 
-			gotURL := sr.GetSafeRedirectURL(tt.redirectURL, tt.fallbackURL)
+			sr := saferedirect.New(tt.allowedHost)
+
+			gotURL := sr.GetSafeRedirectURL(context.Background(), tt.redirectURL, tt.fallbackURL)
 			if gotURL != tt.expectedURL {
 				t.Errorf("GetSafeRedirectURL() = %v, want %v", gotURL, tt.expectedURL)
 			}
@@ -179,9 +184,11 @@ func TestSafeRedirect_GetSafeRedirectURL(t *testing.T) {
 }
 
 func TestSafeRedirect_Redirect(t *testing.T) {
+	t.Parallel()
+
 	tests := []struct {
 		name           string
-		allowedHost    string
+		allowedHost    saferedirect.AllowedHostFunc
 		redirectURL    string
 		fallbackURL    string
 		expectedStatus int
@@ -189,7 +196,7 @@ func TestSafeRedirect_Redirect(t *testing.T) {
 	}{
 		{
 			name:           "safe redirect URL",
-			allowedHost:    "example.com",
+			allowedHost:    saferedirect.StaticHosts("example.com"),
 			redirectURL:    "/dashboard",
 			fallbackURL:    "/home",
 			expectedStatus: http.StatusFound,
@@ -197,7 +204,7 @@ func TestSafeRedirect_Redirect(t *testing.T) {
 		},
 		{
 			name:           "unsafe redirect URL",
-			allowedHost:    "example.com",
+			allowedHost:    saferedirect.StaticHosts("example.com"),
 			redirectURL:    "https://evil.com/phishing",
 			fallbackURL:    "/home",
 			expectedStatus: http.StatusFound,
@@ -205,7 +212,7 @@ func TestSafeRedirect_Redirect(t *testing.T) {
 		},
 		{
 			name:           "empty redirect URL",
-			allowedHost:    "example.com",
+			allowedHost:    saferedirect.StaticHosts("example.com"),
 			redirectURL:    "",
 			fallbackURL:    "/home",
 			expectedStatus: http.StatusFound,
@@ -213,7 +220,7 @@ func TestSafeRedirect_Redirect(t *testing.T) {
 		},
 		{
 			name:           "double slash attack",
-			allowedHost:    "example.com",
+			allowedHost:    saferedirect.StaticHosts("example.com"),
 			redirectURL:    "//evil.com/phishing",
 			fallbackURL:    "/home",
 			expectedStatus: http.StatusFound,
@@ -221,7 +228,7 @@ func TestSafeRedirect_Redirect(t *testing.T) {
 		},
 		{
 			name:           "slash-backslash attack",
-			allowedHost:    "example.com",
+			allowedHost:    saferedirect.StaticHosts("example.com"),
 			redirectURL:    "/\\evil.com/phishing",
 			fallbackURL:    "/home",
 			expectedStatus: http.StatusFound,
@@ -231,22 +238,19 @@ func TestSafeRedirect_Redirect(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			sr := saferedirect.SafeRedirect{
+			t.Parallel()
-				AllowedHost: tt.allowedHost,
+
-			}
+			sr := saferedirect.New(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)
@@ -254,3 +258,129 @@ func TestSafeRedirect_Redirect(t *testing.T) {
 		})
 	}
 }
+
+func TestSafeRedirect_DynamicAllowedHost(t *testing.T) {
+	t.Parallel()
+
+	trustedDomains := map[string]bool{
+		"app.getprobo.com":   true,
+		"trust.company.com":  true,
+		"compliance.acme.io": true,
+	}
+
+	sr := saferedirect.New(func(_ context.Context, host string) bool {
+		return trustedDomains[host]
+	})
+
+	tests := []struct {
+		name        string
+		redirectURL string
+		fallbackURL string
+		expectedURL string
+	}{
+		{
+			name:        "primary host passes",
+			redirectURL: "https://app.getprobo.com/trust/my-slug",
+			fallbackURL: "/",
+			expectedURL: "https://app.getprobo.com/trust/my-slug",
+		},
+		{
+			name:        "trusted custom domain passes",
+			redirectURL: "https://trust.company.com/overview",
+			fallbackURL: "/",
+			expectedURL: "https://trust.company.com/overview",
+		},
+		{
+			name:        "another trusted custom domain passes",
+			redirectURL: "https://compliance.acme.io/documents",
+			fallbackURL: "/",
+			expectedURL: "https://compliance.acme.io/documents",
+		},
+		{
+			name:        "untrusted domain rejected",
+			redirectURL: "https://evil.com/phishing",
+			fallbackURL: "/",
+			expectedURL: "/",
+		},
+		{
+			name:        "relative path still works",
+			redirectURL: "/trust/my-slug",
+			fallbackURL: "/",
+			expectedURL: "/trust/my-slug",
+		},
+		{
+			name:        "javascript scheme rejected",
+			redirectURL: "javascript:alert('xss')",
+			fallbackURL: "/",
+			expectedURL: "/",
+		},
+		{
+			name:        "empty redirect URL uses fallback",
+			redirectURL: "",
+			fallbackURL: "/",
+			expectedURL: "/",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			w := httptest.NewRecorder()
+			r := httptest.NewRequest("GET", "http://test.com", nil)
+
+			sr.Redirect(w, r, tt.redirectURL, tt.fallbackURL, http.StatusFound)
+
+			location := w.Header().Get("Location")
+			if location != tt.expectedURL {
+				t.Errorf("Redirect() location = %v, want %v", location, tt.expectedURL)
+			}
+		})
+	}
+}
+
+func TestStaticHosts(t *testing.T) {
+	t.Parallel()
+
+	t.Run("single host", func(t *testing.T) {
+		t.Parallel()
+
+		fn := saferedirect.StaticHosts("example.com")
+		if !fn(context.Background(), "example.com") {
+			t.Error("expected example.com to be allowed")
+		}
+		if fn(context.Background(), "other.com") {
+			t.Error("expected other.com to be rejected")
+		}
+	})
+
+	t.Run("multiple hosts", func(t *testing.T) {
+		t.Parallel()
+
+		fn := saferedirect.StaticHosts("a.com", "b.com", "c.com")
+		if !fn(context.Background(), "a.com") {
+			t.Error("expected a.com to be allowed")
+		}
+		if !fn(context.Background(), "b.com") {
+			t.Error("expected b.com to be allowed")
+		}
+		if !fn(context.Background(), "c.com") {
+			t.Error("expected c.com to be allowed")
+		}
+		if fn(context.Background(), "d.com") {
+			t.Error("expected d.com to be rejected")
+		}
+	})
+
+	t.Run("empty strings ignored", func(t *testing.T) {
+		t.Parallel()
+
+		fn := saferedirect.StaticHosts("", "example.com")
+		if fn(context.Background(), "") {
+			t.Error("expected empty host to be rejected")
+		}
+		if !fn(context.Background(), "example.com") {
+			t.Error("expected example.com to be allowed")
+		}
+	})
+}

pkg/server/api/api.go

@@ -15,6 +15,7 @@
 package api
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net/http"
@@ -188,6 +189,14 @@ func NewServer(cfg Config) (*Server, error) {
 			cfg.Cookie,
 			cfg.TokenSecret,
 			cfg.BaseURL,
+			func(ctx context.Context, host string) bool {
+				if host == cfg.BaseURL.Host() {
+					return true
+				}
+
+				_, err := cfg.Trust.GetByDomainName(ctx, host)
+				return err == nil
+			},
 		),
 	}, nil
 }

pkg/server/api/connect/v1/oidc_handler.go

@@ -22,7 +22,6 @@ import (
 	"github.com/go-chi/chi/v5"
 	"go.gearno.de/kit/httpserver"
 	"go.gearno.de/kit/log"
-	"go.probo.inc/probo/pkg/baseurl"
 	"go.probo.inc/probo/pkg/coredata"
 	"go.probo.inc/probo/pkg/iam"
 	"go.probo.inc/probo/pkg/saferedirect"
@@ -33,18 +32,21 @@ import (
 type OIDCHandler struct {
 	iam           *iam.Service
 	sessionCookie *authn.Cookie
-	baseURL       *baseurl.BaseURL
 	logger        *log.Logger
 	safeRedirect  *saferedirect.SafeRedirect
 }
 
-func NewOIDCHandler(iam *iam.Service, cookieConfig securecookie.Config, baseURL *baseurl.BaseURL, logger *log.Logger) *OIDCHandler {
+func NewOIDCHandler(
+	iam *iam.Service,
+	cookieConfig securecookie.Config,
+	logger *log.Logger,
+	allowedHost saferedirect.AllowedHostFunc,
+) *OIDCHandler {
 	return &OIDCHandler{
 		iam:           iam,
 		sessionCookie: authn.NewCookie(&cookieConfig),
-		baseURL:       baseURL,
 		logger:        logger,
-		safeRedirect:  &saferedirect.SafeRedirect{AllowedHost: baseURL.Host()},
+		safeRedirect:  saferedirect.New(allowedHost),
 	}
 }
 

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

@@ -39,6 +39,7 @@ import (
 	"go.probo.inc/probo/pkg/baseurl"
 	"go.probo.inc/probo/pkg/gid"
 	"go.probo.inc/probo/pkg/iam"
+	"go.probo.inc/probo/pkg/saferedirect"
 	"go.probo.inc/probo/pkg/securecookie"
 	"go.probo.inc/probo/pkg/server/api/authn"
 	"go.probo.inc/probo/pkg/server/api/authz"
@@ -55,7 +56,14 @@ type (
 	}
 )
 
-func NewMux(logger *log.Logger, svc *iam.Service, cookieConfig securecookie.Config, tokenSecret string, baseURL *baseurl.BaseURL) *chi.Mux {
+func NewMux(
+	logger *log.Logger,
+	svc *iam.Service,
+	cookieConfig securecookie.Config,
+	tokenSecret string,
+	baseURL *baseurl.BaseURL,
+	allowedRedirectHost saferedirect.AllowedHostFunc,
+) *chi.Mux {
 	r := chi.NewMux()
 
 	sessionMiddleware := authn.NewSessionMiddleware(svc, cookieConfig)
@@ -66,7 +74,7 @@ func NewMux(logger *log.Logger, svc *iam.Service, cookieConfig securecookie.Conf
 
 	router := r.With(sessionMiddleware, apiKeyMiddleware)
 
-	oidcHandler := NewOIDCHandler(svc, cookieConfig, baseURL, logger)
+	oidcHandler := NewOIDCHandler(svc, cookieConfig, logger, allowedRedirectHost)
 
 	router.Handle("/graphql", graphqlHandler)
 	router.Get("/saml/2.0/metadata", samlHandler.MetadataHandler)

pkg/server/api/connect/v1/saml_handler.go

@@ -45,7 +45,7 @@ func NewSAMLHandler(iam *iam.Service, cookieConfig securecookie.Config, baseURL
 		sessionCookie: authn.NewCookie(&cookieConfig),
 		baseURL:       baseURL,
 		logger:        logger,
-		safeRedirect:  &saferedirect.SafeRedirect{AllowedHost: baseURL.Host()},
+		safeRedirect:  saferedirect.New(saferedirect.StaticHosts(baseURL.Host())),
 	}
 }
 

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

@@ -81,7 +81,7 @@ func NewMux(
 ) *chi.Mux {
 	r := chi.NewMux()
 
-	safeRedirect := &saferedirect.SafeRedirect{AllowedHost: baseURL.Host()}
+	safeRedirect := saferedirect.New(saferedirect.StaticHosts(baseURL.Host()))
 
 	graphqlHandler := NewGraphQLHandler(iamSvc, proboSvc, esignSvc, mailmanSvc, customDomainCname, logger)
 
@@ -141,9 +141,9 @@ func NewMux(
 			var oauthSafeRedirect *saferedirect.SafeRedirect
 			switch provider {
 			case "SLACK":
-				oauthSafeRedirect = &saferedirect.SafeRedirect{AllowedHost: "slack.com"}
+				oauthSafeRedirect = saferedirect.New(saferedirect.StaticHosts("slack.com"))
 			case "GOOGLE_WORKSPACE":
-				oauthSafeRedirect = &saferedirect.SafeRedirect{AllowedHost: "accounts.google.com"}
+				oauthSafeRedirect = saferedirect.New(saferedirect.StaticHosts("accounts.google.com"))
 			}
 			oauthSafeRedirect.Redirect(w, r, redirectURL, "/", http.StatusSeeOther)
 		})

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

@@ -8778,6 +8778,7 @@ func (r *riskConnectionResolver) TotalCount(ctx context.Context, obj *types.Risk
 func (r *slackConnectionResolver) Permission(ctx context.Context, obj *types.SlackConnection, action string) (bool, error) {
 	return r.Resolver.Permission(ctx, obj, action)
 }
+
 // Organization is the resolver for the organization field.
 func (r *snapshotResolver) Organization(ctx context.Context, obj *types.Snapshot) (*types.Organization, error) {
 	if err := r.authorize(ctx, obj.ID, probo.ActionOrganizationGet); err != nil {
@@ -10716,6 +10717,7 @@ func (r *Resolver) RiskConnection() schema.RiskConnectionResolver { return &risk
 func (r *Resolver) SlackConnection() schema.SlackConnectionResolver {
 	return &slackConnectionResolver{r}
 }
+
 // Snapshot returns schema.SnapshotResolver implementation.
 func (r *Resolver) Snapshot() schema.SnapshotResolver { return &snapshotResolver{r} }
 

pkg/server/api/trust/v1/v1_resolver.go

@@ -193,10 +193,10 @@ func (r *mutationResolver) SendMagicLink(ctx context.Context, input types.SendMa
 
 	baseURL := compliancepage.CompliancePageBaseURLFromContext(ctx)
 
-	safeRedirect := &saferedirect.SafeRedirect{AllowedHost: baseurl.MustParse(*baseURL).Host()}
+	safeRedirect := saferedirect.New(saferedirect.StaticHosts(baseurl.MustParse(*baseURL).Host()))
 
 	if input.Continue != nil {
-		_, ok := safeRedirect.Validate(*input.Continue)
+		_, ok := safeRedirect.Validate(ctx, *input.Continue)
 		if !ok {
 			return nil, gqlutils.Invalidf(ctx, "invalid continue URL")
 		}