ory-am/hydra

// Copyright © 2022 Ory Corp
// SPDX-License-Identifier: Apache-2.0

package driver

import (
    "context"
    "io/fs"
    "strings"
    "time"

    "github.com/gobuffalo/pop/v6"
    _ "github.com/jackc/pgx/v4/stdlib"
    "github.com/luna-duclos/instrumentedsql"

    "github.com/ory/hydra/v2/client"
    "github.com/ory/hydra/v2/consent"
    "github.com/ory/hydra/v2/hsm"
    "github.com/ory/hydra/v2/jwk"
    "github.com/ory/hydra/v2/oauth2/trust"
    "github.com/ory/hydra/v2/persistence/sql"
    "github.com/ory/hydra/v2/x"
    "github.com/ory/x/contextx"
    "github.com/ory/x/dbal"
    "github.com/ory/x/errorsx"
    otelsql "github.com/ory/x/otelx/sql"
    "github.com/ory/x/popx"
    "github.com/ory/x/resilience"
    "github.com/ory/x/sqlcon"
)

type RegistrySQL struct {
    *RegistryBase
    defaultKeyManager jwk.Manager
    initialPing       func(r *RegistrySQL) error
}

var _ Registry = new(RegistrySQL)

// defaultInitialPing is the default function that will be called within RegistrySQL.Init to make sure
// the database is reachable. It can be injected for test purposes by changing the value
// of RegistrySQL.initialPing.
var defaultInitialPing = func(m *RegistrySQL) error {
    if err := resilience.Retry(m.l, 5*time.Second, 5*time.Minute, m.Ping); err != nil {
        m.Logger().Print("Could not ping database: ", err)
        return errorsx.WithStack(err)
    }
    return nil
}

func init() {
    dbal.RegisterDriver(
        func() dbal.Driver {
            return NewRegistrySQL()
        },
    )
}

func NewRegistrySQL() *RegistrySQL {
    r := &RegistrySQL{
        RegistryBase: new(RegistryBase),
        initialPing:  defaultInitialPing,
    }
    r.RegistryBase.with(r)
    return r
}

func (m *RegistrySQL) Init(
    ctx context.Context,
    skipNetworkInit bool,
    migrate bool,
    ctxer contextx.Contextualizer,
    extraMigrations []fs.FS,
    goMigrations []popx.Migration,
) error {
    if m.persister == nil {
        m.WithContextualizer(ctxer)
        var opts []instrumentedsql.Opt
        if m.Tracer(ctx).IsLoaded() {
            opts = []instrumentedsql.Opt{
                instrumentedsql.WithTracer(otelsql.NewTracer()),
                instrumentedsql.WithOmitArgs(), // don't risk leaking PII or secrets
                instrumentedsql.WithOpsExcluded(instrumentedsql.OpSQLRowsNext),
            }
        }

        // new db connection
        pool, idlePool, connMaxLifetime, connMaxIdleTime, cleanedDSN := sqlcon.ParseConnectionOptions(
            m.l, m.Config().DSN(),
        )
        c, err := pop.NewConnection(
            &pop.ConnectionDetails{
                URL:                       sqlcon.FinalizeDSN(m.l, cleanedDSN),
                IdlePool:                  idlePool,
                ConnMaxLifetime:           connMaxLifetime,
                ConnMaxIdleTime:           connMaxIdleTime,
                Pool:                      pool,
                UseInstrumentedDriver:     m.Tracer(ctx).IsLoaded(),
                InstrumentedDriverOptions: opts,
                Unsafe:                    m.Config().DbIgnoreUnknownTableColumns(),
            },
        )
        if err != nil {
            return errorsx.WithStack(err)
        }
        if err := resilience.Retry(m.l, 5*time.Second, 5*time.Minute, c.Open); err != nil {
            return errorsx.WithStack(err)
        }

        p, err := sql.NewPersister(ctx, c, m, m.Config(), extraMigrations, goMigrations)
        if err != nil {
            return err
        }
        m.persister = p
        if err := m.initialPing(m); err != nil {
            return err
        }

        if m.Config().HSMEnabled() {
            hardwareKeyManager := hsm.NewKeyManager(m.HSMContext(), m.Config())
            m.defaultKeyManager = jwk.NewManagerStrategy(hardwareKeyManager, m.persister)
        } else {
            m.defaultKeyManager = m.persister
        }

        // if dsn is memory we have to run the migrations on every start
        // use case - such as
        // - just in memory
        // - shared connection
        // - shared but unique in the same process
        // see: https://sqlite.org/inmemorydb.html
        if dbal.IsMemorySQLite(m.Config().DSN()) {
            m.Logger().Print("Hydra is running migrations on every startup as DSN is memory.\n")
            m.Logger().Print("This means your data is lost when Hydra terminates.\n")
            if err := p.MigrateUp(context.Background()); err != nil {
                return err
            }
        } else if migrate {
            if err := p.MigrateUp(context.Background()); err != nil {
                return err
            }
        }

        if skipNetworkInit {
            m.persister = p
        } else {
            net, err := p.DetermineNetwork(ctx)
            if err != nil {
                m.Logger().WithError(err).Warnf("Unable to determine network, retrying.")
                return err
            }

            m.persister = p.WithFallbackNetworkID(net.ID)
        }

        if m.Config().HSMEnabled() {
            hardwareKeyManager := hsm.NewKeyManager(m.HSMContext(), m.Config())
            m.defaultKeyManager = jwk.NewManagerStrategy(hardwareKeyManager, m.persister)
        } else {
            m.defaultKeyManager = m.persister
        }

    }

    return nil
}

func (m *RegistrySQL) alwaysCanHandle(dsn string) bool {
    scheme := strings.Split(dsn, "://")[0]
    s := dbal.Canonicalize(scheme)
    return s == dbal.DriverMySQL || s == dbal.DriverPostgreSQL || s == dbal.DriverCockroachDB
}

func (m *RegistrySQL) Ping() error {
    return m.Persister().Ping()
}

func (m *RegistrySQL) ClientManager() client.Manager {
    return m.Persister()
}

func (m *RegistrySQL) ConsentManager() consent.Manager {
    return m.Persister()
}

func (m *RegistrySQL) OAuth2Storage() x.FositeStorer {
    return m.Persister()
}

func (m *RegistrySQL) KeyManager() jwk.Manager {
    return m.defaultKeyManager
}

func (m *RegistrySQL) SoftwareKeyManager() jwk.Manager {
    return m.Persister()
}

func (m *RegistrySQL) GrantManager() trust.GrantManager {
    return m.Persister()
}