diff --git a/internal/store/natsstore/keys.go b/internal/store/natsstore/keys.go new file mode 100644 index 00000000000..ad790e343d4 --- /dev/null +++ b/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 +} diff --git a/internal/store/natsstore/keys_test.go b/internal/store/natsstore/keys_test.go new file mode 100644 index 00000000000..1d4c27a0abc --- /dev/null +++ b/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) + } + }) +} diff --git a/internal/store/natsstore/kv.go b/internal/store/natsstore/kv.go new file mode 100644 index 00000000000..aea3712fb74 --- /dev/null +++ b/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()} +} diff --git a/internal/store/natsstore/kv_test.go b/internal/store/natsstore/kv_test.go new file mode 100644 index 00000000000..50dd6922615 --- /dev/null +++ b/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) + } +}