From d99789dfbd17ab0aa96e0d5af61b8206c256ee68 Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Sat, 7 Nov 2026 17:27:20 -0500 Subject: [PATCH] feat(natsstore): bootstrap the store buckets idempotently Operator-provisioned clusters must be verified rather than modified, and replicas start concurrently, so bucket creation is opt-in and safe to race. Signed-off-by: Yordis Prieto --- internal/store/natsstore/bootstrap.go | 173 +++++++++++++++++++ internal/store/natsstore/bootstrap_test.go | 185 +++++++++++++++++++++ internal/store/natsstore/buckets.go | 173 +++++++++++++++++++ internal/store/natsstore/buckets_test.go | 80 +++++++++ 4 files changed, 611 insertions(+) create mode 100644 internal/store/natsstore/bootstrap.go create mode 100644 internal/store/natsstore/bootstrap_test.go create mode 100644 internal/store/natsstore/buckets.go create mode 100644 internal/store/natsstore/buckets_test.go diff --git a/internal/store/natsstore/bootstrap.go b/internal/store/natsstore/bootstrap.go new file mode 100644 index 00000000000..92059222d75 --- /dev/null +++ b/internal/store/natsstore/bootstrap.go @@ -0,0 +1,173 @@ +package natsstore + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "time" + + "github.com/nats-io/nats.go" + "github.com/nats-io/nats.go/jetstream" + log "github.com/sirupsen/logrus" +) + +const ( + // SchemaVersion is the layout version this binary reads and writes. + SchemaVersion = 1 + // SchemaKey holds the schema marker in the meta bucket. + SchemaKey = "schema" +) + +// ErrSchemaTooNew is returned when the stored layout was written by a newer binary. +var ErrSchemaTooNew = errors.New("stored schema is newer than this binary supports") + +// ErrBucketMissing is returned in CreateNever mode when a bucket has not been provisioned. +var ErrBucketMissing = errors.New("bucket does not exist") + +// DriftError reports a bucket whose configuration differs from the expected shape. +type DriftError struct { + Bucket string + // Fields lists each mismatch as `name: have X, want Y`. + Fields []string +} + +func (e *DriftError) Error() string { + return fmt.Sprintf("bucket %s differs from the expected configuration (%s); update it to match or set %s=%s to let startup reconcile it", + e.Bucket, strings.Join(e.Fields, "; "), EnvStoreCreate, CreateAuto) +} + +// Bootstrapped is a connected JetStream context with every bucket verified. +type Bootstrapped struct { + JS jetstream.JetStream + Features ServerFeatures + // Buckets maps full bucket name to its handle. + Buckets map[string]jetstream.KeyValue +} + +// Bucket returns the handle for a purpose such as BucketAuth. +func (b *Bootstrapped) Bucket(settings StoreSettings, purpose string) jetstream.KeyValue { + return b.Buckets[settings.BucketName(purpose)] +} + +// Bootstrap probes the server, ensures every bucket and records the schema +// marker. It is idempotent and safe for many replicas to run concurrently. +func Bootstrap(ctx context.Context, nc *nats.Conn, settings StoreSettings) (*Bootstrapped, error) { + if err := settings.Validate(); err != nil { + return nil, err + } + js, err := jetstream.New(nc) + if err != nil { + return nil, fmt.Errorf("failed to create JetStream context: %w", err) + } + features, err := Probe(ctx, nc, js) + if err != nil { + return nil, err + } + if !features.KeyTTL { + return nil, fmt.Errorf("NATS server %s does not support per-key TTL (JetStream API level %d); use a server of version %d.%d or newer", + features.Version, features.APILevel, 2, 11) + } + if !features.AtLeast(TargetMajor, TargetMinor) { + log.Warnf("NATS server %s is older than the tested %d.%d line", features.Version, TargetMajor, TargetMinor) + } + + buckets := make(map[string]jetstream.KeyValue) + for _, spec := range settings.Buckets() { + kv, err := EnsureBucket(ctx, js, spec, settings.Create) + if err != nil { + return nil, err + } + buckets[spec.Name] = kv + } + meta := buckets[settings.BucketName(BucketMeta)] + if err := EnsureSchema(ctx, meta); err != nil { + return nil, err + } + return &Bootstrapped{JS: js, Features: features, Buckets: buckets}, nil +} + +// EnsureBucket makes the bucket exist with the expected shape. With CreateAuto +// it creates or reconciles; with CreateNever it only looks. Either way the +// final shape is verified, so a drifted bucket is reported instead of used. +func EnsureBucket(ctx context.Context, js jetstream.JetStream, spec BucketSpec, mode CreateMode) (jetstream.KeyValue, error) { + var ( + kv jetstream.KeyValue + err error + ) + switch mode { + case CreateAuto: + kv, err = js.CreateOrUpdateKeyValue(ctx, spec.config()) + case CreateNever: + kv, err = js.KeyValue(ctx, spec.Name) + if errors.Is(err, jetstream.ErrBucketNotFound) { + return nil, fmt.Errorf("%w: %s; provision it or set %s=%s", ErrBucketMissing, spec.Name, EnvStoreCreate, CreateAuto) + } + default: + return nil, fmt.Errorf("unknown create mode %q", mode) + } + if err != nil { + return nil, fmt.Errorf("failed to ensure bucket %s: %w", spec.Name, err) + } + if err := verifyBucket(ctx, kv, spec); err != nil { + return nil, err + } + return kv, nil +} + +func verifyBucket(ctx context.Context, kv jetstream.KeyValue, spec BucketSpec) error { + status, err := kv.Status(ctx) + if err != nil { + return fmt.Errorf("failed to read bucket %s: %w", spec.Name, err) + } + cfg := status.(*jetstream.KeyValueBucketStatus).StreamInfo().Config + + var drift []string + check := func(name string, have, want any) { + if have != want { + drift = append(drift, fmt.Sprintf("%s: have %v, want %v", name, have, want)) + } + } + check("storage", cfg.Storage.String(), spec.Storage.String()) + check("replicas", cfg.Replicas, spec.Replicas) + check("history", cfg.MaxMsgsPerSubject, int64(spec.History)) + check("ttl", cfg.MaxAge, time.Duration(0)) + check("limit marker ttl", cfg.SubjectDeleteMarkerTTL, spec.MarkerTTL) + if len(drift) > 0 { + return &DriftError{Bucket: spec.Name, Fields: drift} + } + return nil +} + +type schemaMarker struct { + Version int `json:"v"` +} + +// EnsureSchema records the schema version once and refuses a newer stored one, +// so an old binary in a rolling upgrade never writes into a layout it does not know. +func EnsureSchema(ctx context.Context, meta jetstream.KeyValue) error { + payload, err := json.Marshal(schemaMarker{Version: SchemaVersion}) + if err != nil { + return err + } + _, err = meta.Create(ctx, SchemaKey, payload) + if err == nil { + return nil + } + if !errors.Is(err, jetstream.ErrKeyExists) { + return fmt.Errorf("failed to record schema marker: %w", err) + } + entry, err := meta.Get(ctx, SchemaKey) + if err != nil { + return fmt.Errorf("failed to read schema marker: %w", err) + } + var stored schemaMarker + if err := json.Unmarshal(entry.Value(), &stored); err != nil { + return fmt.Errorf("schema marker is not valid: %w", err) + } + if stored.Version > SchemaVersion { + return fmt.Errorf("%w: stored v%d, supported v%d", ErrSchemaTooNew, stored.Version, SchemaVersion) + } + return nil +} diff --git a/internal/store/natsstore/bootstrap_test.go b/internal/store/natsstore/bootstrap_test.go new file mode 100644 index 00000000000..27add32fc98 --- /dev/null +++ b/internal/store/natsstore/bootstrap_test.go @@ -0,0 +1,185 @@ +package natsstore + +import ( + "errors" + "strings" + "sync" + "testing" + "time" + + "github.com/nats-io/nats.go/jetstream" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/store/natsstore/natstest" +) + +func autoSettings(prefix string) StoreSettings { + settings := DefaultStoreSettings() + settings.Prefix = prefix + settings.Create = CreateAuto + return settings +} + +func TestBootstrapIsIdempotentAndConcurrencySafe(t *testing.T) { + srv := natstest.Start(t) + settings := autoSettings("idem") + ctx := testContext(t) + + var wg sync.WaitGroup + errs := make(chan error, 6) + for range 6 { + wg.Add(1) + go func() { + defer wg.Done() + nc, err := Connect(ctx, SettingsFromEnv(mapEnv(map[string]string{EnvURL: srv.URL})).Resolve(), OSLookup, 5*time.Second) + if err != nil { + errs <- err + return + } + defer nc.Close() + _, err = Bootstrap(ctx, nc, settings) + errs <- err + }() + } + wg.Wait() + close(errs) + for err := range errs { + if err != nil { + t.Fatalf("concurrent bootstrap: %v", err) + } + } + + nc, err := connect(t, map[string]string{EnvURL: srv.URL}, 5*time.Second) + if err != nil { + t.Fatal(err) + } + boot, err := Bootstrap(ctx, nc, settings) + if err != nil { + t.Fatalf("repeat bootstrap: %v", err) + } + if len(boot.Buckets) != len(settings.Buckets()) { + t.Fatalf("buckets = %d", len(boot.Buckets)) + } + entry, err := boot.Bucket(settings, BucketMeta).Get(ctx, SchemaKey) + if err != nil { + t.Fatal(err) + } + if string(entry.Value()) != `{"v":1}` { + t.Fatalf("schema marker = %s", entry.Value()) + } + if entry.Revision() != 1 { + t.Fatalf("schema marker was rewritten, revision %d", entry.Revision()) + } +} + +func TestBootstrapNeverModeVerifiesOnly(t *testing.T) { + srv := natstest.Start(t) + nc, err := connect(t, map[string]string{EnvURL: srv.URL}, 5*time.Second) + if err != nil { + t.Fatal(err) + } + ctx := testContext(t) + verify := DefaultStoreSettings() + verify.Prefix = "verify" + + _, err = Bootstrap(ctx, nc, verify) + if !errors.Is(err, ErrBucketMissing) { + t.Fatalf("err = %v, want ErrBucketMissing", err) + } + js, _ := jetstream.New(nc) + if _, lookupErr := js.KeyValue(ctx, "verify-config"); !errors.Is(lookupErr, jetstream.ErrBucketNotFound) { + t.Fatalf("never mode created a bucket: %v", lookupErr) + } + + if _, err := Bootstrap(ctx, nc, autoSettings("verify")); err != nil { + t.Fatal(err) + } + if _, err := Bootstrap(ctx, nc, verify); err != nil { + t.Fatalf("verify after provisioning: %v", err) + } +} + +func TestBootstrapReportsDrift(t *testing.T) { + srv := natstest.Start(t) + nc, err := connect(t, map[string]string{EnvURL: srv.URL}, 5*time.Second) + if err != nil { + t.Fatal(err) + } + ctx := testContext(t) + js, _ := jetstream.New(nc) + if _, err := js.CreateKeyValue(ctx, jetstream.KeyValueConfig{Bucket: "drift-config", History: 3}); err != nil { + t.Fatal(err) + } + + settings := DefaultStoreSettings() + settings.Prefix = "drift" + _, err = Bootstrap(ctx, nc, settings) + var drift *DriftError + if !errors.As(err, &drift) { + t.Fatalf("err = %v, want DriftError", err) + } + if drift.Bucket != "drift-config" || len(drift.Fields) != 1 || !strings.HasPrefix(drift.Fields[0], "history") { + t.Fatalf("drift = %+v", drift) + } +} + +func TestBootstrapRefusesNewerSchema(t *testing.T) { + srv := natstest.Start(t) + nc, err := connect(t, map[string]string{EnvURL: srv.URL}, 5*time.Second) + if err != nil { + t.Fatal(err) + } + ctx := testContext(t) + settings := autoSettings("schema") + boot, err := Bootstrap(ctx, nc, settings) + if err != nil { + t.Fatal(err) + } + if _, err := boot.Bucket(settings, BucketMeta).Put(ctx, SchemaKey, []byte(`{"v":2}`)); err != nil { + t.Fatal(err) + } + if _, err := Bootstrap(ctx, nc, settings); !errors.Is(err, ErrSchemaTooNew) { + t.Fatalf("err = %v, want ErrSchemaTooNew", err) + } +} + +func TestPerKeyTTLExpiresInBootstrappedBuckets(t *testing.T) { + srv := natstest.Start(t) + nc, err := connect(t, map[string]string{EnvURL: srv.URL}, 5*time.Second) + if err != nil { + t.Fatal(err) + } + ctx := testContext(t) + settings := autoSettings("ttl") + boot, err := Bootstrap(ctx, nc, settings) + if err != nil { + t.Fatal(err) + } + + lease := boot.Bucket(settings, BucketLease) + if _, err := lease.Create(ctx, "short", []byte("x"), jetstream.KeyTTL(time.Second)); err != nil { + t.Fatalf("create with key TTL: %v", err) + } + if _, err := lease.Create(ctx, "short", []byte("y")); !errors.Is(err, jetstream.ErrKeyExists) { + t.Fatalf("second create = %v, want ErrKeyExists", err) + } + deadline := time.Now().Add(15 * time.Second) + for { + _, err := lease.Get(ctx, "short") + if errors.Is(err, jetstream.ErrKeyNotFound) { + break + } + if time.Now().After(deadline) { + t.Fatalf("key still readable after its TTL: %v", err) + } + time.Sleep(250 * time.Millisecond) + } + if _, err := lease.Create(ctx, "short", []byte("again")); err != nil { + t.Fatalf("create after expiry: %v", err) + } + + // Keys in the refresh journal bucket never expire: the bucket has no limit marker TTL. + journal := boot.Bucket(settings, BucketRefresh) + if _, err := journal.Create(ctx, "claim", []byte("x"), jetstream.KeyTTL(time.Second)); err == nil { + t.Fatal("refresh bucket accepted a per-key TTL") + } +} diff --git a/internal/store/natsstore/buckets.go b/internal/store/natsstore/buckets.go new file mode 100644 index 00000000000..4d5f2aa4d7b --- /dev/null +++ b/internal/store/natsstore/buckets.go @@ -0,0 +1,173 @@ +package natsstore + +import ( + "fmt" + "strconv" + "strings" + "time" + + "github.com/nats-io/nats.go/jetstream" +) + +// Store environment variables. +const ( + EnvStorePrefix = "NATS_STORE_PREFIX" + EnvStoreCreate = "NATS_STORE_CREATE" + EnvStoreReplicas = "NATS_STORE_REPLICAS" + EnvStoreStorage = "NATS_STORE_STORAGE" + + // DefaultPrefix names the buckets when NATS_STORE_PREFIX is unset. + DefaultPrefix = "cliproxy" + + maxPrefixLength = 48 + maxReplicas = 5 + + // limitMarkerTTL is how long buckets with per-key TTL keep expiry markers. + limitMarkerTTL = time.Minute +) + +// CreateMode decides whether bootstrap may create or update buckets. +type CreateMode string + +const ( + // CreateNever only verifies that buckets exist with the expected shape, for + // deployments where an operator provisions them with tighter permissions. + CreateNever CreateMode = "never" + // CreateAuto creates missing buckets and updates drifted ones. + CreateAuto CreateMode = "auto" +) + +// StoreSettings is the non-secret store layout configuration. +type StoreSettings struct { + // Prefix names every bucket as -, so deployments can share a cluster. + Prefix string + Create CreateMode + // Replicas is the replication factor, 1 to 5. + Replicas int + // Storage applies to buckets that tolerate loss. Buckets holding credentials + // or the refresh journal always use file storage. + Storage jetstream.StorageType +} + +// DefaultStoreSettings verifies only, with one replica and file storage. +func DefaultStoreSettings() StoreSettings { + return StoreSettings{Prefix: DefaultPrefix, Create: CreateNever, Replicas: 1, Storage: jetstream.FileStorage} +} + +// StoreSettingsFromEnv reads the store layout from the environment, applying defaults for unset values. +func StoreSettingsFromEnv(env Lookup) (StoreSettings, error) { + settings := DefaultStoreSettings() + if value, ok := env(EnvStorePrefix); ok && value != "" { + settings.Prefix = value + } + if value, ok := env(EnvStoreCreate); ok && value != "" { + settings.Create = CreateMode(strings.ToLower(strings.TrimSpace(value))) + } + if value, ok := env(EnvStoreReplicas); ok && value != "" { + n, err := strconv.Atoi(strings.TrimSpace(value)) + if err != nil { + return StoreSettings{}, &SettingError{Setting: EnvStoreReplicas, Err: fmt.Errorf("%q is not a number", value)} + } + settings.Replicas = n + } + if value, ok := env(EnvStoreStorage); ok && value != "" { + switch strings.ToLower(strings.TrimSpace(value)) { + case "file": + settings.Storage = jetstream.FileStorage + case "memory": + settings.Storage = jetstream.MemoryStorage + default: + return StoreSettings{}, &SettingError{Setting: EnvStoreStorage, Err: fmt.Errorf("%q is not file or memory", value)} + } + } + if err := settings.Validate(); err != nil { + return StoreSettings{}, err + } + return settings, nil +} + +// Validate rejects settings that cannot name valid buckets or streams. +func (s StoreSettings) Validate() error { + if s.Prefix == "" || len(s.Prefix) > maxPrefixLength { + return &SettingError{Setting: EnvStorePrefix, Err: fmt.Errorf("must be 1 to %d characters", maxPrefixLength)} + } + for _, r := range s.Prefix { + if !(r >= 'a' && r <= 'z') && !(r >= 'A' && r <= 'Z') && !(r >= '0' && r <= '9') && r != '_' && r != '-' { + return &SettingError{Setting: EnvStorePrefix, Err: fmt.Errorf("allows only letters, digits, - and _")} + } + } + switch s.Create { + case CreateNever, CreateAuto: + default: + return &SettingError{Setting: EnvStoreCreate, Err: fmt.Errorf("%q is not auto or never", s.Create)} + } + if s.Replicas < 1 || s.Replicas > maxReplicas { + return &SettingError{Setting: EnvStoreReplicas, Err: fmt.Errorf("must be 1 to %d", maxReplicas)} + } + return nil +} + +// BucketSpec is the shape one KV bucket must have. +type BucketSpec struct { + // Name is the full bucket name, including the prefix. + Name string + Description string + History uint8 + // MarkerTTL enables per-key TTL on the bucket, zero for none. It is never + // combined with a bucket-wide TTL, so keys only expire when written with one. + MarkerTTL time.Duration + Storage jetstream.StorageType + Replicas int +} + +// Bucket purposes, combined with the prefix to form bucket names. +const ( + BucketConfig = "config" + BucketAuth = "auth" + BucketRefresh = "refresh" + BucketCooldown = "cooldown" + BucketLease = "lease" + BucketMeta = "meta" +) + +// BucketName returns the bucket name for a purpose. +func (s StoreSettings) BucketName(purpose string) string { return s.Prefix + "-" + purpose } + +// Buckets lists every bucket the backend uses, in a stable order. Config, +// credentials, the refresh journal and metadata never expire and always use +// file storage; cooldown and lease keys expire per key and follow Storage. +func (s StoreSettings) Buckets() []BucketSpec { + spec := func(purpose, description string, history uint8, markerTTL time.Duration, durable bool) BucketSpec { + storage := jetstream.FileStorage + if !durable { + storage = s.Storage + } + return BucketSpec{ + Name: s.BucketName(purpose), + Description: description, + History: history, + MarkerTTL: markerTTL, + Storage: storage, + Replicas: s.Replicas, + } + } + return []BucketSpec{ + spec(BucketConfig, "CLIProxyAPI configuration", 20, 0, true), + spec(BucketAuth, "CLIProxyAPI credentials", 5, 0, true), + spec(BucketRefresh, "CLIProxyAPI refresh grant journal", 1, 0, true), + spec(BucketCooldown, "CLIProxyAPI cooldown state", 1, limitMarkerTTL, false), + spec(BucketLease, "CLIProxyAPI refresh leases", 1, limitMarkerTTL, false), + spec(BucketMeta, "CLIProxyAPI schema and migration markers", 1, limitMarkerTTL, true), + } +} + +func (b BucketSpec) config() jetstream.KeyValueConfig { + return jetstream.KeyValueConfig{ + Bucket: b.Name, + Description: b.Description, + History: b.History, + Storage: b.Storage, + Replicas: b.Replicas, + LimitMarkerTTL: b.MarkerTTL, + } +} diff --git a/internal/store/natsstore/buckets_test.go b/internal/store/natsstore/buckets_test.go new file mode 100644 index 00000000000..db3a6701894 --- /dev/null +++ b/internal/store/natsstore/buckets_test.go @@ -0,0 +1,80 @@ +package natsstore + +import ( + "errors" + "testing" + + "github.com/nats-io/nats.go/jetstream" +) + +func TestStoreSettingsFromEnv(t *testing.T) { + got, err := StoreSettingsFromEnv(mapEnv(nil)) + if err != nil { + t.Fatal(err) + } + if got != DefaultStoreSettings() { + t.Fatalf("defaults = %+v", got) + } + + got, err = StoreSettingsFromEnv(mapEnv(map[string]string{ + EnvStorePrefix: "prod-eu", + EnvStoreCreate: "AUTO", + EnvStoreReplicas: "3", + EnvStoreStorage: "memory", + })) + if err != nil { + t.Fatal(err) + } + want := StoreSettings{Prefix: "prod-eu", Create: CreateAuto, Replicas: 3, Storage: jetstream.MemoryStorage} + if got != want { + t.Fatalf("got %+v, want %+v", got, want) + } +} + +func TestStoreSettingsRejectsBadValues(t *testing.T) { + tests := map[string]map[string]string{ + EnvStorePrefix + " dot": {EnvStorePrefix: "a.b"}, + EnvStorePrefix + " wild": {EnvStorePrefix: "a>"}, + EnvStorePrefix + " long": {EnvStorePrefix: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"}, + EnvStoreCreate + " word": {EnvStoreCreate: "always"}, + EnvStoreReplicas + " nan": {EnvStoreReplicas: "three"}, + EnvStoreReplicas + " big": {EnvStoreReplicas: "6"}, + EnvStoreReplicas + " 0": {EnvStoreReplicas: "0"}, + EnvStoreStorage + " word": {EnvStoreStorage: "disk"}, + } + for name, env := range tests { + t.Run(name, func(t *testing.T) { + _, err := StoreSettingsFromEnv(mapEnv(env)) + var setting *SettingError + if !errors.As(err, &setting) { + t.Fatalf("err = %v, want SettingError", err) + } + }) + } +} + +func TestBucketsAreStableAndDurableOnesUseFileStorage(t *testing.T) { + settings := DefaultStoreSettings() + settings.Storage = jetstream.MemoryStorage + buckets := settings.Buckets() + + wantNames := []string{"cliproxy-config", "cliproxy-auth", "cliproxy-refresh", "cliproxy-cooldown", "cliproxy-lease", "cliproxy-meta"} + for i, name := range wantNames { + if buckets[i].Name != name { + t.Fatalf("bucket %d = %q, want %q", i, buckets[i].Name, name) + } + } + for _, bucket := range buckets { + durable := bucket.Name != "cliproxy-cooldown" && bucket.Name != "cliproxy-lease" + if durable && bucket.Storage != jetstream.FileStorage { + t.Errorf("%s uses %s storage", bucket.Name, bucket.Storage) + } + if !durable && bucket.Storage != jetstream.MemoryStorage { + t.Errorf("%s ignores the storage setting", bucket.Name) + } + } + // The refresh journal is never allowed to expire keys. + if journal := buckets[2]; journal.MarkerTTL != 0 || journal.History != 1 { + t.Errorf("refresh bucket = %+v", journal) + } +}