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/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())),
 	}
 }