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/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")
 		}