276 lines
9.7 KiB
Go
276 lines
9.7 KiB
Go
package db_test
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/url"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
db "git.sendico.io/tech/sxchnge/shared/database"
|
|
dependencies "git.sendico.io/tech/sxchnge/shared/database/dependencies"
|
|
storable "git.sendico.io/tech/sxchnge/shared/database/record"
|
|
"git.sendico.io/tech/sxchnge/shared/database/repository"
|
|
merrors "git.sendico.io/tech/sxchnge/shared/errors"
|
|
"git.sendico.io/tech/sxchnge/shared/model"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/testcontainers/testcontainers-go/modules/mongodb"
|
|
"go.mongodb.org/mongo-driver/v2/bson"
|
|
"go.mongodb.org/mongo-driver/v2/mongo"
|
|
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
type rotationRecord struct {
|
|
storable.Base `bson:",inline"`
|
|
Value string `bson:"value"`
|
|
}
|
|
|
|
func (*rotationRecord) Collection() string { return "rotation" }
|
|
|
|
type leaseServer struct {
|
|
admin *mongo.Database
|
|
issued atomic.Int64
|
|
revoked sync.Map
|
|
outage atomic.Bool
|
|
badPassword atomic.Bool
|
|
}
|
|
|
|
func (s *leaseServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|
if s.outage.Load() {
|
|
w.WriteHeader(503)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
switch r.URL.Path {
|
|
case "/v1/auth/approle/login":
|
|
// A unique token owns each credential in this fake Vault HTTP boundary.
|
|
token := fmt.Sprintf("test-token-%d", s.issued.Load()+1)
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"auth": map[string]any{"client_token": token, "lease_duration": 60}})
|
|
case "/v1/database/creds/bff":
|
|
n := s.issued.Add(1)
|
|
user := fmt.Sprintf("lease%d", n)
|
|
err := s.admin.RunCommand(r.Context(), bson.D{{Key: "createUser", Value: user}, {Key: "pwd", Value: "test-password"}, {Key: "roles", Value: bson.A{bson.M{"role": "readWrite", "db": "rotation"}}}}).Err()
|
|
if err != nil {
|
|
w.WriteHeader(500)
|
|
return
|
|
}
|
|
password := "test-password"
|
|
if s.badPassword.Load() {
|
|
password = "wrong-password"
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"lease_id": user, "lease_duration": 8, "data": map[string]string{"username": user, "password": password}})
|
|
case "/v1/auth/token/revoke-self":
|
|
user := "lease" + strings.TrimPrefix(r.Header.Get("X-Vault-Token"), "test-token-")
|
|
_ = s.admin.RunCommand(r.Context(), bson.D{{Key: "dropUser", Value: user}}).Err()
|
|
s.revoked.Store(user, true)
|
|
w.WriteHeader(204)
|
|
default:
|
|
w.WriteHeader(404)
|
|
}
|
|
}
|
|
|
|
func TestRotationPreservesRequestsCursorsAndTransactions(t *testing.T) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
|
|
defer cancel()
|
|
container, err := mongodb.Run(ctx, "mongo:latest", mongodb.WithReplicaSet("rs"), mongodb.WithUsername("root"), mongodb.WithPassword("test-admin-password"))
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { _ = container.Terminate(context.Background()) })
|
|
raw, err := container.ConnectionString(ctx)
|
|
require.NoError(t, err)
|
|
uri, err := url.Parse(raw)
|
|
require.NoError(t, err)
|
|
q := uri.Query()
|
|
q.Set("directConnection", "true")
|
|
q.Set("authSource", "admin")
|
|
uri.RawQuery = q.Encode()
|
|
admin, err := mongo.Connect(options.Client().ApplyURI(uri.String()))
|
|
require.NoError(t, err)
|
|
defer admin.Disconnect(context.Background())
|
|
fake := &leaseServer{admin: admin.Database("admin")}
|
|
vault := httptest.NewServer(fake)
|
|
defer vault.Close()
|
|
directory := t.TempDir()
|
|
role := filepath.Join(directory, "role")
|
|
secret := filepath.Join(directory, "secret")
|
|
require.NoError(t, os.WriteFile(role, []byte("test-role"), 0600))
|
|
require.NoError(t, os.WriteFile(secret, []byte("test-secret"), 0600))
|
|
uri.User = nil
|
|
uri.Path = "/rotation"
|
|
factory, err := db.NewConnection(zap.NewNop(), &db.Config{Driver: db.Mongo, Settings: model.SettingsT{"database": "rotation", "uri": uri.String(), "vault": map[string]any{"address": vault.URL, "role": "bff", "role_id_file": role, "secret_id_file": secret}}})
|
|
require.NoError(t, err)
|
|
defer factory.CloseConnection()
|
|
deps, err := dependencies.ResolveDependencies(factory)
|
|
require.NoError(t, err)
|
|
repo := repository.CreateRepository(deps.Database, "rotation")
|
|
require.NoError(t, repo.Insert(ctx, &rotationRecord{Value: "cursor-fixture"}, nil))
|
|
|
|
var failures, requests atomic.Int64
|
|
httpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if _, e := repo.CountByFilter(r.Context(), repository.Query()); e != nil {
|
|
failures.Add(1)
|
|
w.WriteHeader(503)
|
|
return
|
|
}
|
|
requests.Add(1)
|
|
w.WriteHeader(200)
|
|
}))
|
|
defer httpServer.Close()
|
|
loadCtx, stopLoad := context.WithCancel(ctx)
|
|
var load sync.WaitGroup
|
|
defer func() { stopLoad(); load.Wait() }()
|
|
for range 4 {
|
|
load.Add(1)
|
|
go func() {
|
|
defer load.Done()
|
|
for loadCtx.Err() == nil {
|
|
res, e := http.Get(httpServer.URL)
|
|
if e != nil {
|
|
failures.Add(1)
|
|
} else {
|
|
res.Body.Close()
|
|
if res.StatusCode != 200 {
|
|
failures.Add(1)
|
|
}
|
|
}
|
|
time.Sleep(5 * time.Millisecond)
|
|
}
|
|
}()
|
|
}
|
|
|
|
// A cursor and a transaction must keep generation 1 alive through publication.
|
|
first, scope, release, err := deps.Database.Acquire(ctx)
|
|
require.NoError(t, err)
|
|
firstClient := first.Client()
|
|
release()
|
|
_, _, _, err = deps.Database.Acquire(scope)
|
|
require.Error(t, err, "escaped operation scope must be rejected")
|
|
txEntered := make(chan struct{})
|
|
finishTx := make(chan struct{})
|
|
txDone := make(chan error, 1)
|
|
var entered sync.Once
|
|
go func() {
|
|
_, e := factory.TransactionFactory().CreateTransaction().Execute(ctx, func(tx context.Context) (any, error) {
|
|
if e := repo.Insert(tx, &rotationRecord{Value: "transaction-before"}, nil); e != nil {
|
|
return nil, e
|
|
}
|
|
entered.Do(func() { close(txEntered) })
|
|
select {
|
|
case <-finishTx:
|
|
case <-tx.Done():
|
|
return nil, tx.Err()
|
|
}
|
|
// Nested transaction must share the same session/generation.
|
|
return repo.WithTransaction(tx, func(nested context.Context) (any, error) {
|
|
return nil, repo.Insert(nested, &rotationRecord{Value: "transaction-after"}, nil)
|
|
})
|
|
})
|
|
txDone <- e
|
|
}()
|
|
select {
|
|
case <-txEntered:
|
|
case err := <-txDone:
|
|
t.Fatalf("transaction did not enter: %v", err)
|
|
case <-ctx.Done():
|
|
t.Fatal(ctx.Err())
|
|
}
|
|
cursorEntered := make(chan struct{})
|
|
finishCursor := make(chan struct{})
|
|
cursorDone := make(chan error, 1)
|
|
go func() {
|
|
var once sync.Once
|
|
cursorDone <- repo.FindManyByFilter(ctx, repository.Query(), func(d *mongo.Cursor) error {
|
|
once.Do(func() { close(cursorEntered) })
|
|
select {
|
|
case <-finishCursor:
|
|
return nil
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
}
|
|
})
|
|
}()
|
|
select {
|
|
case <-cursorEntered:
|
|
case err := <-cursorDone:
|
|
t.Fatalf("cursor did not enter: %v", err)
|
|
case <-ctx.Done():
|
|
t.Fatal(ctx.Err())
|
|
}
|
|
require.Eventually(t, func() bool {
|
|
database, _, free, e := deps.Database.Acquire(ctx)
|
|
if e != nil {
|
|
return false
|
|
}
|
|
defer free()
|
|
return database.Client() != firstClient
|
|
}, 6*time.Second, 20*time.Millisecond)
|
|
_, revoked := fake.revoked.Load("lease1")
|
|
require.False(t, revoked, "in-flight work lost its credentials")
|
|
close(finishTx)
|
|
require.NoError(t, <-txDone)
|
|
_, revoked = fake.revoked.Load("lease1")
|
|
require.False(t, revoked, "cursor must retain old pool")
|
|
close(finishCursor)
|
|
require.NoError(t, <-cursorDone)
|
|
require.Eventually(t, func() bool { _, ok := fake.revoked.Load("lease1"); return ok }, 2*time.Second, 20*time.Millisecond)
|
|
count, err := repo.CountByFilter(ctx, repository.Filter("value", "transaction-after"))
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, 1, count)
|
|
|
|
// Reject a candidate with invalid credentials while retaining the valid pool.
|
|
fake.badPassword.Store(true)
|
|
before := fake.issued.Load()
|
|
require.Eventually(t, func() bool { return fake.issued.Load() > before }, 5*time.Second, 20*time.Millisecond)
|
|
require.Eventually(t, func() bool { _, ok := fake.revoked.Load(fmt.Sprintf("lease%d", before+1)); return ok }, time.Second, 10*time.Millisecond)
|
|
require.NoError(t, repo.Insert(ctx, &rotationRecord{Value: "during-bad-refresh"}, nil))
|
|
fake.badPassword.Store(false)
|
|
fake.outage.Store(true)
|
|
require.NoError(t, repo.Insert(ctx, &rotationRecord{Value: "during-vault-outage"}, nil))
|
|
time.Sleep(300 * time.Millisecond)
|
|
fake.outage.Store(false)
|
|
require.Eventually(t, func() bool { return fake.issued.Load() > before+1 }, 3*time.Second, 20*time.Millisecond)
|
|
stopLoad()
|
|
load.Wait()
|
|
require.Greater(t, requests.Load(), int64(100))
|
|
require.Zero(t, failures.Load(), "HTTP requests failed during rotation")
|
|
|
|
// If Vault stays unavailable beyond the lease, fail closed without restarting.
|
|
fake.outage.Store(true)
|
|
require.Eventually(t, func() bool { _, e := repo.CountByFilter(ctx, repository.Query()); return e != nil }, 10*time.Second, 50*time.Millisecond)
|
|
fake.outage.Store(false)
|
|
require.Eventually(t, func() bool { _, e := repo.CountByFilter(ctx, repository.Query()); return e == nil }, 6*time.Second, 20*time.Millisecond)
|
|
// Transaction errors still roll back all writes after generation changes.
|
|
_, err = repo.WithTransaction(ctx, func(tx context.Context) (any, error) {
|
|
if e := repo.Insert(tx, &rotationRecord{Value: "rollback"}, nil); e != nil {
|
|
return nil, e
|
|
}
|
|
return nil, merrors.Internal("injected rollback")
|
|
})
|
|
require.Error(t, err)
|
|
count, err = repo.CountByFilter(ctx, repository.Filter("value", "rollback"))
|
|
require.NoError(t, err)
|
|
require.Zero(t, count)
|
|
held, heldContext, releaseHeld, err := deps.Database.Acquire(ctx)
|
|
require.NoError(t, err)
|
|
manager := deps.Database.(interface{ Close(context.Context) error })
|
|
short, stop := context.WithTimeout(ctx, 30*time.Millisecond)
|
|
require.ErrorIs(t, manager.Close(short), context.DeadlineExceeded)
|
|
stop()
|
|
_, _, _, err = deps.Database.Acquire(ctx)
|
|
require.Error(t, err)
|
|
require.NoError(t, held.Client().Ping(heldContext, nil), "shutdown must drain, not kill held work")
|
|
releaseHeld()
|
|
closed, stop := context.WithTimeout(ctx, 5*time.Second)
|
|
defer stop()
|
|
require.NoError(t, manager.Close(closed))
|
|
|
|
}
|