feat(natsstore): encode keys and wrap buckets with revision CAS

internal/store/natsstore/keys.go

@@ -0,0 +1,36 @@
+package natsstore
+
+import (
+	"encoding/base64"
+	"fmt"
+)
+
+// maxKeyLength bounds an encoded key. A key becomes the last token of a
+// subject, so it stays far below the server control line limit.
+const maxKeyLength = 512
+
+// keyEncoding is unpadded base64url. Its alphabet is a subset of what KV keys
+// allow and has no dot, so an encoded key is always one subject token.
+var keyEncoding = base64.RawURLEncoding
+
+// EncodeKey turns an identifier, which may hold slashes, dots, @ or any UTF-8,
+// into a KV key. The mapping is reversible and collision free.
+func EncodeKey(id string) (string, error) {
+	if id == "" {
+		return "", fmt.Errorf("key id is empty")
+	}
+	key := keyEncoding.EncodeToString([]byte(id))
+	if len(key) > maxKeyLength {
+		return "", fmt.Errorf("key id is too long: encoded length %d exceeds %d", len(key), maxKeyLength)
+	}
+	return key, nil
+}
+
+// DecodeKey returns the identifier a key was encoded from.
+func DecodeKey(key string) (string, error) {
+	raw, err := keyEncoding.DecodeString(key)
+	if err != nil || len(raw) == 0 {
+		return "", fmt.Errorf("key %q is not an encoded identifier", key)
+	}
+	return string(raw), nil
+}

internal/store/natsstore/keys_test.go

@@ -0,0 +1,73 @@
+package natsstore
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestEncodeKeyRoundTrips(t *testing.T) {
+	ids := []string{
+		"a.json",
+		"sub/dir/foo@bar.com.json",
+		"codex-user@example.com-plus.json",
+		"nested/../odd name.json",
+		"ünïcödé/日本.json",
+		".hidden",
+		"trailing.",
+		"x",
+	}
+	seen := map[string]string{}
+	for _, id := range ids {
+		key, err := EncodeKey(id)
+		if err != nil {
+			t.Fatalf("EncodeKey(%q): %v", id, err)
+		}
+		if strings.ContainsAny(key, ".*> =/") {
+			t.Fatalf("key %q for %q is not one safe token", key, id)
+		}
+		if other, dup := seen[key]; dup {
+			t.Fatalf("%q and %q collide on %q", id, other, key)
+		}
+		seen[key] = id
+		back, err := DecodeKey(key)
+		if err != nil || back != id {
+			t.Fatalf("DecodeKey(%q) = %q, %v; want %q", key, back, err, id)
+		}
+	}
+}
+
+func TestEncodeKeyRejectsEmptyAndOversized(t *testing.T) {
+	if _, err := EncodeKey(""); err == nil {
+		t.Fatal("empty id accepted")
+	}
+	if _, err := EncodeKey(strings.Repeat("a", 400)); err == nil {
+		t.Fatal("oversized id accepted")
+	}
+	if _, err := EncodeKey(strings.Repeat("a", 380)); err != nil {
+		t.Fatalf("id at the limit rejected: %v", err)
+	}
+}
+
+func TestDecodeKeyRejectsForeignKeys(t *testing.T) {
+	for _, key := range []string{"", "not base64!", "a.b", "config/main"} {
+		if _, err := DecodeKey(key); err == nil {
+			t.Fatalf("DecodeKey(%q) accepted a foreign key", key)
+		}
+	}
+}
+
+func FuzzEncodeKey(f *testing.F) {
+	for _, seed := range []string{"a", "x/y@z.json", "\xff\xfe", "日本"} {
+		f.Add(seed)
+	}
+	f.Fuzz(func(t *testing.T, id string) {
+		key, err := EncodeKey(id)
+		if err != nil {
+			return
+		}
+		back, err := DecodeKey(key)
+		if err != nil || back != id {
+			t.Fatalf("round trip of %q gave %q, %v", id, back, err)
+		}
+	})
+}

internal/store/natsstore/kv.go

@@ -0,0 +1,142 @@
+package natsstore
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/nats-io/nats.go/jetstream"
+)
+
+var (
+	// ErrNotFound is returned when a key has no live value.
+	ErrNotFound = errors.New("key not found")
+	// ErrConflict is matched by a ConflictError.
+	ErrConflict = errors.New("revision conflict")
+)
+
+// ConflictError reports that a write lost a compare-and-swap, because the key
+// was created, changed or deleted since the revision the caller read.
+type ConflictError struct {
+	Bucket string
+	ID     string
+}
+
+func (e *ConflictError) Error() string {
+	return fmt.Sprintf("revision conflict on %s in bucket %s", e.ID, e.Bucket)
+}
+
+// Is makes errors.Is(err, ErrConflict) match.
+func (e *ConflictError) Is(target error) bool { return target == ErrConflict }
+
+// Entry is one live value with the revision a later write must present.
+type Entry struct {
+	ID       string
+	Value    []byte
+	Revision uint64
+	// Modified is when the revision was written.
+	Modified time.Time
+}
+
+// KV wraps a bucket so callers work with identifiers and revisions and never
+// with encoded keys. Every write is a compare-and-swap; there is no blind put.
+type KV struct {
+	bucket string
+	kv     jetstream.KeyValue
+}
+
+// NewKV wraps the bucket handle.
+func NewKV(bucket string, kv jetstream.KeyValue) *KV { return &KV{bucket: bucket, kv: kv} }
+
+// Get returns the live value for the identifier, or ErrNotFound.
+func (k *KV) Get(ctx context.Context, id string) (Entry, error) {
+	key, err := EncodeKey(id)
+	if err != nil {
+		return Entry{}, err
+	}
+	entry, err := k.kv.Get(ctx, key)
+	if err != nil {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
+			return Entry{}, ErrNotFound
+		}
+		return Entry{}, fmt.Errorf("failed to read %s from %s: %w", id, k.bucket, err)
+	}
+	return entryOf(id, entry), nil
+}
+
+// Create writes a value only if the key has no live value.
+func (k *KV) Create(ctx context.Context, id string, value []byte) (uint64, error) {
+	key, err := EncodeKey(id)
+	if err != nil {
+		return 0, err
+	}
+	rev, err := k.kv.Create(ctx, key, value)
+	if err != nil {
+		return 0, k.writeError("create", id, err)
+	}
+	return rev, nil
+}
+
+// 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)
+	if err != nil {
+		return 0, err
+	}
+	rev, err := k.kv.Update(ctx, key, value, expected)
+	if err != nil {
+		return 0, k.writeError("update", id, err)
+	}
+	return rev, nil
+}
+
+// Delete removes the key only if it is still at the expected revision.
+func (k *KV) Delete(ctx context.Context, id string, expected uint64) error {
+	key, err := EncodeKey(id)
+	if err != nil {
+		return err
+	}
+	if err := k.kv.Delete(ctx, key, jetstream.LastRevision(expected)); err != nil {
+		return k.writeError("delete", id, err)
+	}
+	return nil
+}
+
+// List returns every live entry in key order. Keys deleted while listing are skipped.
+func (k *KV) List(ctx context.Context) ([]Entry, error) {
+	keys, err := k.kv.ListKeys(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to list keys of %s: %w", k.bucket, err)
+	}
+	defer func() { _ = keys.Stop() }()
+
+	var entries []Entry
+	for key := range keys.Keys() {
+		id, err := DecodeKey(key)
+		if err != nil {
+			// Not written by this backend; leave it alone.
+			continue
+		}
+		entry, err := k.kv.Get(ctx, key)
+		if err != nil {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
+				continue
+			}
+			return nil, fmt.Errorf("failed to read %s from %s: %w", id, k.bucket, err)
+		}
+		entries = append(entries, entryOf(id, entry))
+	}
+	return entries, nil
+}
+
+func (k *KV) writeError(op, id string, err error) error {
+	if errors.Is(err, jetstream.ErrKeyExists) {
+		return &ConflictError{Bucket: k.bucket, ID: id}
+	}
+	return fmt.Errorf("failed to %s %s in %s: %w", op, id, k.bucket, err)
+}
+
+func entryOf(id string, entry jetstream.KeyValueEntry) Entry {
+	return Entry{ID: id, Value: entry.Value(), Revision: entry.Revision(), Modified: entry.Created()}
+}

internal/store/natsstore/kv_test.go

@@ -0,0 +1,105 @@
+package natsstore
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/router-for-me/CLIProxyAPI/v8/internal/store/natsstore/natstest"
+)
+
+func openAuthKV(t *testing.T) *KV {
+	t.Helper()
+	srv := natstest.Start(t)
+	nc, err := connect(t, map[string]string{EnvURL: srv.URL}, 5*time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	settings := autoSettings("kvtest")
+	boot, err := Bootstrap(testContext(t), nc, settings)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return NewKV(settings.BucketName(BucketAuth), boot.Bucket(settings, BucketAuth))
+}
+
+func TestKVCompareAndSwap(t *testing.T) {
+	kv := openAuthKV(t)
+	ctx := testContext(t)
+	const id = "sub/dir/a@b.com.json"
+
+	if _, err := kv.Get(ctx, id); !errors.Is(err, ErrNotFound) {
+		t.Fatalf("Get missing = %v, want ErrNotFound", err)
+	}
+	rev1, err := kv.Create(ctx, id, []byte(`{"v":1}`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := kv.Create(ctx, id, []byte(`{"v":2}`)); !errors.Is(err, ErrConflict) {
+		t.Fatalf("second Create = %v, want conflict", err)
+	}
+	rev2, err := kv.Update(ctx, id, []byte(`{"v":2}`), rev1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = kv.Update(ctx, id, []byte(`{"v":3}`), rev1)
+	var conflict *ConflictError
+	if !errors.As(err, &conflict) || conflict.ID != id || conflict.Bucket != "kvtest-auth" {
+		t.Fatalf("stale Update = %v, want ConflictError for %s", err, id)
+	}
+	if err := kv.Delete(ctx, id, rev1); !errors.Is(err, ErrConflict) {
+		t.Fatalf("stale Delete = %v, want conflict", err)
+	}
+	got, err := kv.Get(ctx, id)
+	if err != nil || string(got.Value) != `{"v":2}` || got.Revision != rev2 || got.ID != id {
+		t.Fatalf("Get = %+v, %v", got, err)
+	}
+	if err := kv.Delete(ctx, id, rev2); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := kv.Get(ctx, id); !errors.Is(err, ErrNotFound) {
+		t.Fatalf("Get after delete = %v, want ErrNotFound", err)
+	}
+	if _, err := kv.Update(ctx, id, []byte(`{}`), rev2); !errors.Is(err, ErrConflict) {
+		t.Fatalf("Update after delete = %v, want conflict", err)
+	}
+	if _, err := kv.Create(ctx, id, []byte(`{"v":4}`)); err != nil {
+		t.Fatalf("Create over a delete marker: %v", err)
+	}
+}
+
+func TestKVListSkipsDeletedAndForeignKeys(t *testing.T) {
+	kv := openAuthKV(t)
+	ctx := testContext(t)
+
+	empty, err := kv.List(ctx)
+	if err != nil || len(empty) != 0 {
+		t.Fatalf("List empty = %v, %v", empty, err)
+	}
+	ids := []string{"b.json", "a/c.json", "gone.json"}
+	revs := map[string]uint64{}
+	for _, id := range ids {
+		rev, err := kv.Create(ctx, id, []byte(id))
+		if err != nil {
+			t.Fatal(err)
+		}
+		revs[id] = rev
+	}
+	if err := kv.Delete(ctx, "gone.json", revs["gone.json"]); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := kv.kv.Create(ctx, "plain", []byte("foreign")); err != nil {
+		t.Fatal(err)
+	}
+	entries, err := kv.List(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	got := map[string]string{}
+	for _, entry := range entries {
+		got[entry.ID] = string(entry.Value)
+	}
+	if len(got) != 2 || got["b.json"] != "b.json" || got["a/c.json"] != "a/c.json" {
+		t.Fatalf("List = %v", got)
+	}
+}