From 53fc419fc750d1d251bd06abe6965bed5461988b Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Sat, 7 Nov 2026 21:15:40 -0500 Subject: [PATCH] feat(natsstore): journal refresh grants in a bucket without expiry Replicas sharing credentials must not send the same single-use refresh grant twice. A claim is an atomic create on the grant digest and the receipt a revision CAS over it, so exactly one replica wins and an unfinished claim stays pending forever instead of being reclaimed, which keeps refresh_outcome_unknown meaning what it means for the file journal. The claim treats a delete marker as an existing write, because a plain KV create retries over a marker and would let a deleted claim be taken again. Signed-off-by: Yordis Prieto --- internal/store/natsstore/kv.go | 8 + internal/store/natsstore/refresh_journal.go | 98 ++++++++++ .../store/natsstore/refresh_journal_test.go | 171 ++++++++++++++++++ 3 files changed, 277 insertions(+) create mode 100644 internal/store/natsstore/refresh_journal.go create mode 100644 internal/store/natsstore/refresh_journal_test.go diff --git a/internal/store/natsstore/kv.go b/internal/store/natsstore/kv.go index d4d81545896..3c98f9bf859 100644 --- a/internal/store/natsstore/kv.go +++ b/internal/store/natsstore/kv.go @@ -118,6 +118,14 @@ func (k *KV) Create(ctx context.Context, id string, value []byte) (uint64, error return rev, nil } +// CreateOnce writes a value only if the key was never written. Unlike Create it +// treats a delete marker as a write, so deleting a key does not reopen it. +func (k *KV) CreateOnce(ctx context.Context, id string, value []byte) (uint64, error) { + // An expected revision of zero asks the server for a subject with no + // message at all, markers included, where Create retries over a marker. + return k.Update(ctx, id, value, 0) +} + // Update writes a value only if the key is still at the expected revision. func (k *KV) Update(ctx context.Context, id string, value []byte, expected uint64) (uint64, error) { key, err := EncodeKey(id) diff --git a/internal/store/natsstore/refresh_journal.go b/internal/store/natsstore/refresh_journal.go new file mode 100644 index 00000000000..ba8d7c8293d --- /dev/null +++ b/internal/store/natsstore/refresh_journal.go @@ -0,0 +1,98 @@ +package natsstore + +import ( + "bytes" + "context" + "errors" + "fmt" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" +) + +// RefreshJournal keeps refresh grant claims and receipts in the refresh bucket. +// A claim is an atomic Create, so exactly one replica wins a grant, and the +// receipt replaces the pending marker with a revision compare-and-swap. Nothing +// here deletes, purges or expires an entry: an incomplete claim is never +// reclaimed, so a crash between claim and receipt reads back as pending and the +// grant is not replayed, exactly as with the file journal. +type RefreshJournal struct { + kv *KV +} + +var _ codex.RefreshJournal = (*RefreshJournal)(nil) + +// NewRefreshJournal wraps the refresh bucket. The bucket must have no TTL, as +// BucketSpec describes; The startup probe checks that. +func NewRefreshJournal(kv *KV) *RefreshJournal { return &RefreshJournal{kv: kv} } + +// validRefreshKey accepts only the sha256 hex digest codex.refreshGrantKey +// produces, so a malformed key cannot name another entry. +func validRefreshKey(key string) error { + if len(key) != 64 { + return fmt.Errorf("invalid refresh journal key") + } + for _, char := range key { + if !(char >= '0' && char <= '9' || char >= 'a' && char <= 'f') { + return fmt.Errorf("invalid refresh journal key") + } + } + return nil +} + +// RefreshClaim records the grant as consumed. It reports false when another +// claim, completed or not, already exists. A delete marker counts as a claim, +// so a key someone deleted cannot be claimed again; the grant then reads back +// as an unresolved outcome rather than as unclaimed. +func (j *RefreshJournal) RefreshClaim(ctx context.Context, key string) (bool, error) { + if err := validRefreshKey(key); err != nil { + return false, err + } + if _, err := j.kv.CreateOnce(ctx, key, []byte(codex.PendingRefreshReceipt)); err != nil { + if errors.Is(err, ErrConflict) { + return false, nil + } + return false, err + } + return true, nil +} + +// RefreshRead returns the receipt, the pending marker for a claim without one, +// or nil when the grant was never claimed. +func (j *RefreshJournal) RefreshRead(ctx context.Context, key string) ([]byte, error) { + if err := validRefreshKey(key); err != nil { + return nil, err + } + entry, err := j.kv.Get(ctx, key) + if errors.Is(err, ErrNotFound) { + return nil, nil + } + if err != nil { + return nil, err + } + return entry.Value, nil +} + +// RefreshComplete saves the receipt over the pending claim. It fails when the +// claim is missing or already completed, and never overwrites a receipt. +func (j *RefreshJournal) RefreshComplete(ctx context.Context, key string, receipt []byte) error { + if err := validRefreshKey(key); err != nil { + return err + } + entry, err := j.kv.Get(ctx, key) + if errors.Is(err, ErrNotFound) { + return fmt.Errorf("refresh receipt claim missing") + } + if err != nil { + return err + } + if !bytes.Equal(entry.Value, []byte(codex.PendingRefreshReceipt)) { + return fmt.Errorf("refresh receipt already completed") + } + if _, err := j.kv.Update(ctx, key, receipt, entry.Revision); err != nil { + if errors.Is(err, ErrConflict) { + return fmt.Errorf("refresh receipt already completed") + } + return err + } + return nil +} diff --git a/internal/store/natsstore/refresh_journal_test.go b/internal/store/natsstore/refresh_journal_test.go new file mode 100644 index 00000000000..42cfc3a3c72 --- /dev/null +++ b/internal/store/natsstore/refresh_journal_test.go @@ -0,0 +1,171 @@ +package natsstore + +import ( + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/nats-io/nats.go/jetstream" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" + "github.com/router-for-me/CLIProxyAPI/v8/internal/store/natsstore/natstest" +) + +func openRefreshJournals(t *testing.T, replicas int) (StoreSettings, *Bootstrapped, []*RefreshJournal) { + t.Helper() + srv := natstest.Start(t) + settings := autoSettings("journal") + var journals []*RefreshJournal + var boot *Bootstrapped + for range replicas { + nc, err := connect(t, map[string]string{EnvURL: srv.URL}, 5*time.Second) + if err != nil { + t.Fatal(err) + } + boot, err = Bootstrap(testContext(t), nc, settings) + if err != nil { + t.Fatal(err) + } + journals = append(journals, NewRefreshJournal(NewKV(settings.BucketName(BucketRefresh), boot.Bucket(settings, BucketRefresh)))) + } + return settings, boot, journals +} + +func grantKey(seed string) string { return strings.Repeat(seed, 64)[:64] } + +func TestRefreshJournalClaimReceiptChain(t *testing.T) { + _, _, journals := openRefreshJournals(t, 1) + journal := journals[0] + ctx := testContext(t) + first, second := grantKey("a"), grantKey("b") + + if receipt, err := journal.RefreshRead(ctx, first); err != nil || receipt != nil { + t.Fatalf("read of an unclaimed grant = %q, %v; want nil", receipt, err) + } + if err := journal.RefreshComplete(ctx, first, []byte(`{"value":{}}`)); err == nil { + t.Fatal("completing an unclaimed grant succeeded") + } + if claimed, err := journal.RefreshClaim(ctx, first); err != nil || !claimed { + t.Fatalf("first claim = %v, %v", claimed, err) + } + if receipt, err := journal.RefreshRead(ctx, first); err != nil || string(receipt) != codex.PendingRefreshReceipt { + t.Fatalf("read of a pending claim = %q, %v", receipt, err) + } + // An incomplete claim is never reclaimed, however often it is retried. + for range 3 { + if claimed, err := journal.RefreshClaim(ctx, first); err != nil || claimed { + t.Fatalf("reclaim of a pending claim = %v, %v; want false", claimed, err) + } + } + + rotated := []byte(`{"value":{"access_token":"x","refresh_token":"next"}}`) + if err := journal.RefreshComplete(ctx, first, rotated); err != nil { + t.Fatal(err) + } + if err := journal.RefreshComplete(ctx, first, []byte(`{"rejected":"invalid_grant"}`)); err == nil { + t.Fatal("a receipt was overwritten") + } + if receipt, err := journal.RefreshRead(ctx, first); err != nil || string(receipt) != string(rotated) { + t.Fatalf("read of a completed claim = %q, %v", receipt, err) + } + if claimed, err := journal.RefreshClaim(ctx, first); err != nil || claimed { + t.Fatalf("reclaim of a completed claim = %v, %v; want false", claimed, err) + } + + // The next grant in the chain is an independent claim. + if claimed, err := journal.RefreshClaim(ctx, second); err != nil || !claimed { + t.Fatalf("claim of the rotated grant = %v, %v", claimed, err) + } +} + +func TestRefreshJournalRejectsMalformedKeys(t *testing.T) { + _, _, journals := openRefreshJournals(t, 1) + ctx := testContext(t) + for _, key := range []string{"", "abc", strings.Repeat("A", 64), strings.Repeat("g", 64), strings.Repeat("a", 65)} { + if _, err := journals[0].RefreshClaim(ctx, key); err == nil { + t.Errorf("claim accepted key %q", key) + } + if _, err := journals[0].RefreshRead(ctx, key); err == nil { + t.Errorf("read accepted key %q", key) + } + if err := journals[0].RefreshComplete(ctx, key, []byte("{}")); err == nil { + t.Errorf("complete accepted key %q", key) + } + } +} + +func TestRefreshJournalConcurrentClaimersHaveOneWinner(t *testing.T) { + _, _, journals := openRefreshJournals(t, 4) + ctx := testContext(t) + key := grantKey("c") + + var winners atomic.Int32 + var wg sync.WaitGroup + start := make(chan struct{}) + for _, journal := range journals { + for range 3 { + wg.Add(1) + go func() { + defer wg.Done() + <-start + claimed, err := journal.RefreshClaim(ctx, key) + if err != nil { + t.Errorf("claim: %v", err) + return + } + if claimed { + winners.Add(1) + } + }() + } + } + close(start) + wg.Wait() + if winners.Load() != 1 { + t.Fatalf("claim winners = %d, want 1", winners.Load()) + } + // Every replica sees the same pending claim. + for _, journal := range journals { + if receipt, err := journal.RefreshRead(ctx, key); err != nil || string(receipt) != codex.PendingRefreshReceipt { + t.Fatalf("read = %q, %v", receipt, err) + } + } +} + +func TestRefreshJournalBucketIsDurableWithoutExpiry(t *testing.T) { + settings, boot, _ := openRefreshJournals(t, 1) + status, err := boot.Bucket(settings, BucketRefresh).Status(testContext(t)) + if err != nil { + t.Fatal(err) + } + cfg := status.(*jetstream.KeyValueBucketStatus).StreamInfo().Config + if cfg.Storage != jetstream.FileStorage || cfg.MaxAge != 0 || cfg.SubjectDeleteMarkerTTL != 0 || cfg.Replicas != settings.Replicas { + t.Fatalf("journal bucket = storage %v, max age %v, marker ttl %v, replicas %d", cfg.Storage, cfg.MaxAge, cfg.SubjectDeleteMarkerTTL, cfg.Replicas) + } +} + +func TestRefreshJournalDeletedClaimIsNotReclaimed(t *testing.T) { + settings, boot, journals := openRefreshJournals(t, 1) + ctx := testContext(t) + key := grantKey("d") + if claimed, err := journals[0].RefreshClaim(ctx, key); err != nil || !claimed { + t.Fatalf("claim = %v, %v", claimed, err) + } + // Anything allowed to publish on the bucket can write a delete marker; + // the journal never does. + encoded, err := EncodeKey(key) + if err != nil { + t.Fatal(err) + } + if err := boot.Bucket(settings, BucketRefresh).Delete(ctx, encoded); err != nil { + t.Fatal(err) + } + if claimed, err := journals[0].RefreshClaim(ctx, key); err != nil || claimed { + t.Fatalf("reclaim over a delete marker = %v, %v; want false", claimed, err) + } + if err := journals[0].RefreshComplete(ctx, key, []byte(`{"value":{}}`)); err == nil { + t.Fatal("a receipt was written over a delete marker") + } +}