Keep Mongo operations online across Vault credential rotation
This commit is contained in:
Stephan D committed
1 parent
80d6917aa7
commit
23f7b78c3b
30 files changed
+1290
-202
No files matched your search
@@ -0,0 +1,253 @@
|
||||
package source
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
merrors "git.sendico.io/tech/sxchnge/shared/errors"
|
||||
mlogger "git.sendico.io/tech/sxchnge/shared/logging"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo/readpref"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// Lease owns a credential and its revocation capability. Neither may be logged.
|
||||
type Lease struct {
|
||||
URI string
|
||||
ExpiresAt time.Time
|
||||
Revoke func(context.Context) error
|
||||
}
|
||||
type Issuer interface {
|
||||
Issue(context.Context) (Lease, error)
|
||||
}
|
||||
|
||||
type generation struct {
|
||||
client *mongo.Client
|
||||
lease Lease
|
||||
refs int
|
||||
retired bool
|
||||
drained chan struct{}
|
||||
}
|
||||
type pin struct {
|
||||
manager *Manager
|
||||
generation *generation
|
||||
active atomic.Bool
|
||||
}
|
||||
type pinKey struct{}
|
||||
|
||||
// Manager publishes a verified replacement before retiring the previous pool.
|
||||
// Every cursor, operation and transaction retains its originating generation.
|
||||
type Manager struct {
|
||||
mu sync.Mutex
|
||||
current *generation
|
||||
closed bool
|
||||
issuer Issuer
|
||||
database string
|
||||
logger mlogger.Logger
|
||||
cancel context.CancelFunc
|
||||
workerDone chan struct{}
|
||||
draining sync.WaitGroup
|
||||
}
|
||||
|
||||
func New(ctx context.Context, logger mlogger.Logger, database string, issuer Issuer) (*Manager, error) {
|
||||
if database == "" || issuer == nil {
|
||||
return nil, merrors.InvalidArgument("database and credential issuer are required")
|
||||
}
|
||||
life, cancel := context.WithCancel(context.Background())
|
||||
m := &Manager{issuer: issuer, database: database, logger: logger, cancel: cancel, workerDone: make(chan struct{})}
|
||||
g, err := m.prepare(ctx)
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, err
|
||||
}
|
||||
m.current = g
|
||||
go m.run(life)
|
||||
return m, nil
|
||||
}
|
||||
|
||||
func (m *Manager) Acquire(ctx context.Context) (*mongo.Database, context.Context, func(), error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, ctx, nil, err
|
||||
}
|
||||
if p, ok := ctx.Value(pinKey{}).(*pin); ok && p.manager == m {
|
||||
if !p.active.Load() {
|
||||
return nil, ctx, nil, merrors.FailedPrecondition("database operation scope has ended")
|
||||
}
|
||||
return p.generation.client.Database(m.database), ctx, func() {}, nil
|
||||
}
|
||||
m.mu.Lock()
|
||||
g := m.current
|
||||
if m.closed || g == nil || !time.Now().Before(g.lease.ExpiresAt) {
|
||||
m.mu.Unlock()
|
||||
return nil, ctx, nil, merrors.ServiceUnavailable("no valid database credential is available")
|
||||
}
|
||||
if session := mongo.SessionFromContext(ctx); session != nil && session.Client() != g.client {
|
||||
m.mu.Unlock()
|
||||
return nil, ctx, nil, merrors.FailedPrecondition("transaction has no matching database operation scope")
|
||||
}
|
||||
g.refs++
|
||||
m.mu.Unlock()
|
||||
p := &pin{manager: m, generation: g}
|
||||
p.active.Store(true)
|
||||
bounded, cancel := context.WithDeadline(ctx, g.lease.ExpiresAt)
|
||||
scoped := context.WithValue(bounded, pinKey{}, p)
|
||||
release := func() {
|
||||
if !p.active.CompareAndSwap(true, false) {
|
||||
return
|
||||
}
|
||||
cancel()
|
||||
m.mu.Lock()
|
||||
g.refs--
|
||||
if g.retired && g.refs == 0 {
|
||||
close(g.drained)
|
||||
}
|
||||
m.mu.Unlock()
|
||||
}
|
||||
return g.client.Database(m.database), scoped, release, nil
|
||||
}
|
||||
|
||||
func (m *Manager) prepare(ctx context.Context) (*generation, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
lease, err := m.issuer.Issue(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var client *mongo.Client
|
||||
success := false
|
||||
defer func() {
|
||||
if !success {
|
||||
if client != nil {
|
||||
cleanup, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
_ = client.Disconnect(cleanup)
|
||||
cancel()
|
||||
}
|
||||
m.draining.Add(1)
|
||||
go func() { defer m.draining.Done(); m.revoke(lease) }()
|
||||
}
|
||||
}()
|
||||
if lease.Revoke == nil || time.Until(lease.ExpiresAt) < 2*time.Second {
|
||||
return nil, merrors.FailedPrecondition("database credential lifetime is too short")
|
||||
}
|
||||
client, err = mongo.Connect(options.Client().ApplyURI(lease.URI).SetServerSelectionTimeout(5 * time.Second))
|
||||
if err != nil {
|
||||
return nil, merrors.OpaqueWrap(merrors.ErrServiceUnavailable, err, "create replacement database pool")
|
||||
}
|
||||
if err = client.Ping(ctx, readpref.Primary()); err != nil {
|
||||
return nil, merrors.OpaqueWrap(merrors.ErrServiceUnavailable, err, "validate replacement database pool")
|
||||
}
|
||||
if time.Until(lease.ExpiresAt) < time.Second {
|
||||
return nil, merrors.FailedPrecondition("replacement database credential expired during validation")
|
||||
}
|
||||
success = true
|
||||
return &generation{client: client, lease: lease, drained: make(chan struct{})}, nil
|
||||
}
|
||||
|
||||
func (m *Manager) retireLocked(g *generation) {
|
||||
g.retired = true
|
||||
if g.refs == 0 {
|
||||
close(g.drained)
|
||||
}
|
||||
m.draining.Add(1)
|
||||
go func() {
|
||||
defer m.draining.Done()
|
||||
<-g.drained
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := g.client.Disconnect(ctx); err != nil {
|
||||
m.logger.Warn("Retired database pool close failed")
|
||||
}
|
||||
m.revoke(g.lease)
|
||||
m.logger.Info("Retired database pool drained")
|
||||
}()
|
||||
}
|
||||
|
||||
func (m *Manager) revoke(lease Lease) {
|
||||
if lease.Revoke == nil {
|
||||
return
|
||||
}
|
||||
// Expiration is the backstop when Vault is unavailable during bounded cleanup.
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
defer cancel()
|
||||
if err := lease.Revoke(ctx); err != nil {
|
||||
m.logger.Warn("Database credential cleanup failed; Vault expiry will revoke it")
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func (m *Manager) run(ctx context.Context) {
|
||||
defer close(m.workerDone)
|
||||
m.mu.Lock()
|
||||
wait := time.Until(m.current.lease.ExpiresAt) / 2
|
||||
m.mu.Unlock()
|
||||
backoff := 250 * time.Millisecond
|
||||
failureReported, expiryReported := false, false
|
||||
for {
|
||||
timer := time.NewTimer(wait)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
timer.Stop()
|
||||
return
|
||||
case <-timer.C:
|
||||
}
|
||||
next, err := m.prepare(ctx)
|
||||
if err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
m.mu.Lock()
|
||||
remaining := time.Until(m.current.lease.ExpiresAt)
|
||||
m.mu.Unlock()
|
||||
if !failureReported {
|
||||
m.logger.Warn("Database credential refresh failed; retaining current pool", zap.Duration("valid_for", max(remaining, 0)))
|
||||
failureReported = true
|
||||
}
|
||||
if remaining <= 0 && !expiryReported {
|
||||
m.logger.Error("Database credentials expired; rejecting new operations until refresh succeeds")
|
||||
expiryReported = true
|
||||
}
|
||||
wait = backoff
|
||||
if remaining > 0 {
|
||||
wait = max(100*time.Millisecond, min(wait, remaining/4))
|
||||
}
|
||||
backoff = min(5*time.Second, backoff*2)
|
||||
|
||||
continue
|
||||
}
|
||||
m.mu.Lock()
|
||||
if m.closed {
|
||||
m.retireLocked(next)
|
||||
m.mu.Unlock()
|
||||
return
|
||||
}
|
||||
old := m.current
|
||||
m.current = next
|
||||
m.retireLocked(old)
|
||||
m.mu.Unlock()
|
||||
failureReported, expiryReported = false, false
|
||||
backoff = 250 * time.Millisecond
|
||||
m.logger.Info("Database pool replaced without restarting service")
|
||||
wait = time.Until(next.lease.ExpiresAt) / 2
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Manager) Close(ctx context.Context) error {
|
||||
m.mu.Lock()
|
||||
if !m.closed {
|
||||
m.closed = true
|
||||
m.cancel()
|
||||
m.retireLocked(m.current)
|
||||
}
|
||||
m.mu.Unlock()
|
||||
done := make(chan struct{})
|
||||
go func() { <-m.workerDone; m.draining.Wait(); close(done) }()
|
||||
select {
|
||||
case <-done:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user