chore: reduce registry code bloat and improve DI (#3794)

Author: aeneasr

client/handler.go

@@ -818,13 +818,26 @@ func (h *Handler) ValidDynamicAuth(r *http.Request, ps httprouter.Params) (fosit
 	}
 
 	token := strings.TrimPrefix(fosite.AccessTokenFromRequest(r), "ory_at_")
-	if err := h.r.OAuth2HMACStrategy().Enigma.Validate(r.Context(), token); err != nil {
+	if err := h.r.OAuth2HMACStrategy().ValidateAccessToken(
+		r.Context(),
+		// The strategy checks the expiry time of the token. Registration tokens don't expire (we don't have a way of
+		// rotating them) so we set the expiry time to a time in the future.
+		&fosite.Request{
+			Session: &fosite.DefaultSession{
+				ExpiresAt: map[fosite.TokenType]time.Time{
+					fosite.AccessToken: time.Now().Add(time.Hour),
+				},
+			},
+			RequestedAt: time.Now(),
+		},
+		token,
+	); err != nil {
 		return nil, herodot.ErrUnauthorized.
 			WithTrace(err).
 			WithReason("The requested OAuth 2.0 client does not exist or you provided incorrect credentials.").WithDebug(err.Error())
 	}
 
-	signature := h.r.OAuth2HMACStrategy().Enigma.Signature(token)
+	signature := h.r.OAuth2EnigmaStrategy().Signature(token)
 	if subtle.ConstantTimeCompare([]byte(c.RegistrationAccessTokenSignature), []byte(signature)) == 0 {
 		return nil, errors.WithStack(herodot.ErrUnauthorized.
 			WithReason("The requested OAuth 2.0 client does not exist or you provided incorrect credentials.").WithDebug("Registration access tokens do not match."))

client/registry.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/ory/fosite"
 	foauth2 "github.com/ory/fosite/handler/oauth2"
+	enigma "github.com/ory/fosite/token/hmac"
 	"github.com/ory/hydra/v2/jwk"
 	"github.com/ory/hydra/v2/x"
 )
@@ -22,6 +23,7 @@ type Registry interface {
 	ClientManager() Manager
 	ClientHasher() fosite.Hasher
 	OpenIDJWTStrategy() jwk.JWTSigner
-	OAuth2HMACStrategy() *foauth2.HMACSHAStrategy
+	OAuth2HMACStrategy() foauth2.CoreStrategy
+	OAuth2EnigmaStrategy() *enigma.HMACStrategy
 	config.Provider
 }

consent/strategy_logout_test.go

@@ -40,6 +40,7 @@ func TestLogoutFlows(t *testing.T) {
 	reg := internal.NewMockedRegistry(t, &contextx.Default{})
 	reg.Config().MustSet(ctx, config.KeyAccessTokenStrategy, "opaque")
 	reg.Config().MustSet(ctx, config.KeyConsentRequestMaxAge, time.Hour)
+
 	reg.WithKratos(fakeKratos)
 
 	defaultRedirectedMessage := "redirected to default server"

driver/di.go

@@ -0,0 +1,70 @@
+// Copyright © 2024 Ory Corp
+// SPDX-License-Identifier: Apache-2.0
+
+package driver
+
+import (
+	"github.com/pkg/errors"
+	"go.opentelemetry.io/otel/trace"
+
+	"github.com/ory/fosite"
+	"github.com/ory/fosite/handler/oauth2"
+	"github.com/ory/hydra/v2/consent"
+	"github.com/ory/hydra/v2/driver/config"
+	"github.com/ory/hydra/v2/fositex"
+	"github.com/ory/hydra/v2/hsm"
+	"github.com/ory/hydra/v2/internal/kratos"
+	"github.com/ory/x/contextx"
+	"github.com/ory/x/logrusx"
+)
+
+// WritableRegistry is a deprecated interface that should not be used anymore.
+//
+// Deprecate this at some point.
+type WritableRegistry interface {
+	// WithBuildInfo(v, h, d string) Registry
+
+	WithConfig(c *config.DefaultProvider) Registry
+	WithContextualizer(ctxer contextx.Contextualizer) Registry
+	WithLogger(l *logrusx.Logger) Registry
+	WithTracer(t trace.Tracer) Registry
+	WithTracerWrapper(TracerWrapper) Registry
+	WithKratos(k kratos.Client) Registry
+	WithExtraFositeFactories(f []fositex.Factory) Registry
+	ExtraFositeFactories() []fositex.Factory
+	WithOAuth2Provider(f fosite.OAuth2Provider)
+	WithConsentStrategy(c consent.Strategy)
+	WithHsmContext(h hsm.Context)
+}
+
+type RegistryModifier func(r Registry) error
+
+func WithRegistryModifiers(f ...RegistryModifier) OptionsModifier {
+	return func(o *Options) {
+		o.registryModifiers = f
+	}
+}
+
+func RegistryWithHMACSHAStrategy(s func(r Registry) oauth2.CoreStrategy) RegistryModifier {
+	return func(r Registry) error {
+		switch rt := r.(type) {
+		case *RegistrySQL:
+			rt.hmacs = s(r)
+		default:
+			return errors.Errorf("unable to set HMAC strategy on registry of type %T", r)
+		}
+		return nil
+	}
+}
+
+func RegistryWithHsmContext(h hsm.Context) RegistryModifier {
+	return func(r Registry) error {
+		switch rt := r.(type) {
+		case *RegistrySQL:
+			rt.hsm = h
+		default:
+			return errors.Errorf("unable to set HMAC strategy on registry of type %T", r)
+		}
+		return nil
+	}
+}

driver/factory.go

@@ -19,40 +19,45 @@ import (
 )
 
 type (
-	options struct {
+	Options struct {
 		preload  bool
 		validate bool
 		opts     []configx.OptionModifier
 		config   *config.DefaultProvider
 		// The first default refers to determining the NID at startup; the second default referes to the fact that the Contextualizer may dynamically change the NID.
-		skipNetworkInit  bool
-		tracerWrapper    TracerWrapper
-		extraMigrations  []fs.FS
-		goMigrations     []popx.Migration
-		fositexFactories []fositex.Factory
-		inspect          func(Registry) error
+		skipNetworkInit   bool
+		tracerWrapper     TracerWrapper
+		extraMigrations   []fs.FS
+		goMigrations      []popx.Migration
+		fositexFactories  []fositex.Factory
+		registryModifiers []RegistryModifier
+		inspect           func(Registry) error
 	}
-	OptionsModifier func(*options)
+	OptionsModifier func(*Options)
 
 	TracerWrapper func(*otelx.Tracer) *otelx.Tracer
 )
 
-func newOptions() *options {
-	return &options{
+func NewOptions(opts []OptionsModifier) *Options {
+	o := &Options{
 		validate: true,
 		preload:  true,
 		opts:     []configx.OptionModifier{},
 	}
+	for _, f := range opts {
+		f(o)
+	}
+	return o
 }
 
 func WithConfig(config *config.DefaultProvider) OptionsModifier {
-	return func(o *options) {
+	return func(o *Options) {
 		o.config = config
 	}
 }
 
 func WithOptions(opts ...configx.OptionModifier) OptionsModifier {
-	return func(o *options) {
+	return func(o *Options) {
 		o.opts = append(o.opts, opts...)
 	}
 }
@@ -61,61 +66,58 @@ func WithOptions(opts ...configx.OptionModifier) OptionsModifier {
 //
 // This does not affect schema validation!
 func DisableValidation() OptionsModifier {
-	return func(o *options) {
+	return func(o *Options) {
 		o.validate = false
 	}
 }
 
 // DisablePreloading will not preload the config.
 func DisablePreloading() OptionsModifier {
-	return func(o *options) {
+	return func(o *Options) {
 		o.preload = false
 	}
 }
 
 func SkipNetworkInit() OptionsModifier {
-	return func(o *options) {
+	return func(o *Options) {
 		o.skipNetworkInit = true
 	}
 }
 
 // WithTracerWrapper sets a function that wraps the tracer.
 func WithTracerWrapper(wrapper TracerWrapper) OptionsModifier {
-	return func(o *options) {
+	return func(o *Options) {
 		o.tracerWrapper = wrapper
 	}
 }
 
 // WithExtraMigrations specifies additional database migration.
 func WithExtraMigrations(m ...fs.FS) OptionsModifier {
-	return func(o *options) {
+	return func(o *Options) {
 		o.extraMigrations = append(o.extraMigrations, m...)
 	}
 }
 
 func WithGoMigrations(m ...popx.Migration) OptionsModifier {
-	return func(o *options) {
+	return func(o *Options) {
 		o.goMigrations = append(o.goMigrations, m...)
 	}
 }
 
 func WithExtraFositeFactories(f ...fositex.Factory) OptionsModifier {
-	return func(o *options) {
+	return func(o *Options) {
 		o.fositexFactories = append(o.fositexFactories, f...)
 	}
 }
 
 func Inspect(f func(Registry) error) OptionsModifier {
-	return func(o *options) {
+	return func(o *Options) {
 		o.inspect = f
 	}
 }
 
 func New(ctx context.Context, sl *servicelocatorx.Options, opts []OptionsModifier) (Registry, error) {
-	o := newOptions()
-	for _, f := range opts {
-		f(o)
-	}
+	o := NewOptions(opts)
 
 	l := sl.Logger()
 	if l == nil {
@@ -151,6 +153,12 @@ func New(ctx context.Context, sl *servicelocatorx.Options, opts []OptionsModifie
 
 	r.WithExtraFositeFactories(o.fositexFactories)
 
+	for _, f := range o.registryModifiers {
+		if err := f(r); err != nil {
+			return nil, err
+		}
+	}
+
 	if err = r.Init(ctx, o.skipNetworkInit, false, ctxter, o.extraMigrations, o.goMigrations); err != nil {
 		l.WithError(err).Error("Unable to initialize service registry.")
 		return nil, err

driver/registry.go

@@ -8,23 +8,16 @@ import (
 	"io/fs"
 	"net/http"
 
-	"go.opentelemetry.io/otel/trace"
-
-	"github.com/ory/hydra/v2/fositex"
-	"github.com/ory/hydra/v2/internal/kratos"
-	"github.com/ory/x/httprouterx"
+	enigma "github.com/ory/fosite/token/hmac"
 	"github.com/ory/x/popx"
 
 	"github.com/ory/hydra/v2/aead"
-	"github.com/ory/hydra/v2/hsm"
+	"github.com/ory/hydra/v2/internal/kratos"
 	"github.com/ory/x/contextx"
+	"github.com/ory/x/httprouterx"
 
 	"github.com/ory/hydra/v2/oauth2/trust"
 
-	"github.com/pkg/errors"
-
-	"github.com/ory/x/errorsx"
-
 	"github.com/ory/fosite"
 	foauth2 "github.com/ory/fosite/handler/oauth2"
 
@@ -47,22 +40,13 @@ import (
 
 type Registry interface {
 	dbal.Driver
+	WritableRegistry
 
 	Init(ctx context.Context, skipNetworkInit bool, migrate bool, ctxer contextx.Contextualizer, extraMigrations []fs.FS, goMigrations []popx.Migration) error
 
-	WithBuildInfo(v, h, d string) Registry
-	WithConfig(c *config.DefaultProvider) Registry
-	WithContextualizer(ctxer contextx.Contextualizer) Registry
-	WithLogger(l *logrusx.Logger) Registry
-	WithTracer(t trace.Tracer) Registry
-	WithTracerWrapper(TracerWrapper) Registry
-	WithKratos(k kratos.Client) Registry
 	x.HTTPClientProvider
 	GetJWKSFetcherStrategy() fosite.JWKSFetcherStrategy
 
-	WithExtraFositeFactories(f []fositex.Factory) Registry
-	ExtraFositeFactories() []fositex.Factory
-
 	contextx.Provider
 	config.Provider
 	persistence.Provider
@@ -86,35 +70,37 @@ type Registry interface {
 	ConsentHandler() *consent.Handler
 	OAuth2Handler() *oauth2.Handler
 	HealthHandler() *healthx.Handler
+	OAuth2EnigmaStrategy() *enigma.HMACStrategy
 	OAuth2AwareMiddleware() func(h http.Handler) http.Handler
 
-	OAuth2HMACStrategy() *foauth2.HMACSHAStrategy
-	WithOAuth2Provider(f fosite.OAuth2Provider)
-	WithConsentStrategy(c consent.Strategy)
-	WithHsmContext(h hsm.Context)
+	OAuth2HMACStrategy() foauth2.CoreStrategy
 }
 
 func NewRegistryFromDSN(ctx context.Context, c *config.DefaultProvider, l *logrusx.Logger, skipNetworkInit bool, migrate bool, ctxer contextx.Contextualizer) (Registry, error) {
 	registry, err := NewRegistryWithoutInit(c, l)
 	if err != nil {
 		return nil, err
 	}
+
 	if err := registry.Init(ctx, skipNetworkInit, migrate, ctxer, nil, nil); err != nil {
 		return nil, err
 	}
+
 	return registry, nil
 }
 
 func NewRegistryWithoutInit(c *config.DefaultProvider, l *logrusx.Logger) (Registry, error) {
-	driver, err := dbal.GetDriverFor(c.DSN())
-	if err != nil {
-		return nil, errorsx.WithStack(err)
-	}
-	registry, ok := driver.(Registry)
-	if !ok {
-		return nil, errors.Errorf("driver of type %T does not implement interface Registry", driver)
+	registry := NewRegistrySQL(
+		c, l, config.Version, config.Commit, config.Date,
+	)
+
+	if !registry.CanHandle(c.DSN()) {
+		if dbal.IsSQLite(c.DSN()) {
+			return nil, dbal.ErrSQLiteSupportMissing
+		}
+
+		return nil, dbal.ErrNoResponsibleDriverFound
 	}
-	registry = registry.WithLogger(l).WithConfig(c).WithBuildInfo(config.Version, config.Commit, config.Date)
 
 	return registry, nil
 }