diff --git a/cmd/ateapi/internal/store/ateredis/ateredis.go b/cmd/ateapi/internal/store/ateredis/ateredis.go index 2a27fa7c1..ff1a4770b 100644 --- a/cmd/ateapi/internal/store/ateredis/ateredis.go +++ b/cmd/ateapi/internal/store/ateredis/ateredis.go @@ -15,12 +15,19 @@ // Package ateredis is an ate storage backend built on Redis. // // Actors are stored in keys of the form -// `actor::`. They are -// stored as DBActor JSON-serialized objects, which lets us manipulate them from -// Redis lua. +// `actor::`, holding a binary protobuf-serialized Actor +// record. Binary protobuf roughly halves per-record memory versus protojson and +// is ~4x faster to encode; the optimistic-concurrency version check stays in Go +// (comparing the decoded version field, never the raw bytes), so the encoding is +// never inspected from Redis lua. // // Workers are stored in keys of the form -// `worker:::`, holding a DBWorker JSON object. +// `worker:::`, holding a binary +// protobuf-serialized Worker record. +// +// Only the resident Actor/Worker records use binary protobuf. Atespace records +// and the worker-change pub/sub payloads intentionally stay on protojson: they +// are low-cardinality or ephemeral, so they do not drive RAM at scale. // // Note that redis lua scripting has a restriction that informed the data design // here -- a lua script must predeclare all keys it is going to access. It @@ -311,7 +318,7 @@ func (s *Persistence) GetActor(ctx context.Context, atespace, id string) (*ateap } actor := &ateapipb.Actor{} - if err := protojson.Unmarshal(dbActorBytes, actor); err != nil { + if err := proto.Unmarshal(dbActorBytes, actor); err != nil { return nil, fmt.Errorf("while unmarshaling actor: %w", err) } @@ -330,9 +337,9 @@ func (s *Persistence) CreateActor(ctx context.Context, actor *ateapipb.Actor) er dbActor := proto.Clone(actor).(*ateapipb.Actor) dbActor.Version = 1 - dbActorBytes, err := protojson.Marshal(dbActor) + dbActorBytes, err := proto.Marshal(dbActor) if err != nil { - return fmt.Errorf("in protojson.Marshal: %w", err) + return fmt.Errorf("in proto.Marshal: %w", err) } ok, err := s.rdb.SetNX(ctx, dbKey, dbActorBytes, 0).Result() @@ -354,9 +361,9 @@ func (s *Persistence) CreateWorker(ctx context.Context, worker *ateapipb.Worker) dbWorker := proto.Clone(worker).(*ateapipb.Worker) dbWorker.Version = 1 - dbWorkerBytes, err := protojson.Marshal(dbWorker) + dbWorkerBytes, err := proto.Marshal(dbWorker) if err != nil { - return fmt.Errorf("in protojson.Marshal: %w", err) + return fmt.Errorf("in proto.Marshal: %w", err) } ok, err := s.rdb.SetNX(ctx, dbKey, dbWorkerBytes, 0).Result() @@ -383,8 +390,8 @@ func (s *Persistence) GetWorker(ctx context.Context, namespace, pool, pod string } worker := &ateapipb.Worker{} - if err := protojson.Unmarshal(dbWorkerBytes, worker); err != nil { - return nil, fmt.Errorf("in protojson.Unmarshal: %w", err) + if err := proto.Unmarshal(dbWorkerBytes, worker); err != nil { + return nil, fmt.Errorf("in proto.Unmarshal: %w", err) } if worker.GetWorkerNamespace() != namespace || worker.GetWorkerPool() != pool || worker.GetWorkerPod() != pod { @@ -411,8 +418,8 @@ func (s *Persistence) UpdateWorker(ctx context.Context, worker *ateapipb.Worker, } currentWorker := &ateapipb.Worker{} - if err := protojson.Unmarshal(currentVal, currentWorker); err != nil { - return fmt.Errorf("in protojson.Unmarshal: %w", err) + if err := proto.Unmarshal(currentVal, currentWorker); err != nil { + return fmt.Errorf("in proto.Unmarshal: %w", err) } if currentWorker.GetVersion() != expectedVersion { @@ -432,9 +439,9 @@ func (s *Persistence) UpdateWorker(ctx context.Context, worker *ateapipb.Worker, return fmt.Errorf("ip is immutable") } - newVal, err := protojson.Marshal(dbWorker) + newVal, err := proto.Marshal(dbWorker) if err != nil { - return fmt.Errorf("in protojson.Marshal: %w", err) + return fmt.Errorf("in proto.Marshal: %w", err) } _, err = tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error { @@ -479,8 +486,8 @@ func (s *Persistence) DeleteActor(ctx context.Context, atespace, id string) erro } currentActor := &ateapipb.Actor{} - if err := protojson.Unmarshal(currentVal, currentActor); err != nil { - return fmt.Errorf("in protojson.Unmarshal: %w", err) + if err := proto.Unmarshal(currentVal, currentActor); err != nil { + return fmt.Errorf("in proto.Unmarshal: %w", err) } if currentActor.GetStatus() != ateapipb.Actor_STATUS_SUSPENDED { @@ -521,8 +528,8 @@ func (s *Persistence) UpdateActor(ctx context.Context, actor *ateapipb.Actor, ex } currentActor := &ateapipb.Actor{} - if err := protojson.Unmarshal(currentVal, currentActor); err != nil { - return fmt.Errorf("in protojson.Unmarshal: %w", err) + if err := proto.Unmarshal(currentVal, currentActor); err != nil { + return fmt.Errorf("in proto.Unmarshal: %w", err) } if currentActor.GetVersion() != expectedVersion { @@ -542,9 +549,9 @@ func (s *Persistence) UpdateActor(ctx context.Context, actor *ateapipb.Actor, ex return fmt.Errorf("actor_template_name is immutable") } - newVal, err := protojson.Marshal(dbActor) + newVal, err := proto.Marshal(dbActor) if err != nil { - return fmt.Errorf("in protojson.Marshal: %w", err) + return fmt.Errorf("in proto.Marshal: %w", err) } _, err = tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error { @@ -585,8 +592,16 @@ func (s *Persistence) ListWorkers(ctx context.Context) ([]*ateapipb.Worker, erro } worker := &ateapipb.Worker{} - if err := protojson.Unmarshal([]byte(getCmd.Val()), worker); err != nil { - return fmt.Errorf("in protojson.Unmarshal: %w", err) + if err := proto.Unmarshal([]byte(getCmd.Val()), worker); err != nil { + return fmt.Errorf("in proto.Unmarshal: %w", err) + } + + // Binary protobuf decodes an empty value to a zero message without + // error, so validate the decoded identity against the key. This + // guards against an empty/corrupt key surfacing as a blank worker + // (GetWorker performs the same check). + if worker.GetWorkerNamespace() != parts[1] || worker.GetWorkerPool() != parts[2] || worker.GetWorkerPod() != parts[3] { + return fmt.Errorf("worker record at key %q does not match its key identity", workerKey) } mu.Lock() @@ -754,8 +769,15 @@ func (s *Persistence) fetchActors(ctx context.Context, master *redis.Client, key return nil, fmt.Errorf("while fetching keys in shard %s: %w", master.Options().Addr, err) } + // Pipelined returns one result per queued command, in order, so cmds[i] + // corresponds to keys[i]. We rely on that to validate each decoded record + // against its key below; fail closed if that invariant ever breaks. + if len(cmds) != len(keys) { + return nil, fmt.Errorf("(impossible) pipeline returned %d results for %d keys in shard %s", len(cmds), len(keys), master.Options().Addr) + } + var actors []*ateapipb.Actor - for _, cmd := range cmds { + for i, cmd := range cmds { getCmd, ok := cmd.(*redis.StringCmd) if !ok { continue @@ -768,8 +790,18 @@ func (s *Persistence) fetchActors(ctx context.Context, master *redis.Client, key } actor := &ateapipb.Actor{} - if err := protojson.Unmarshal([]byte(getCmd.Val()), actor); err != nil { - return nil, fmt.Errorf("in protojson.Unmarshal: %w", err) + if err := proto.Unmarshal([]byte(getCmd.Val()), actor); err != nil { + return nil, fmt.Errorf("in proto.Unmarshal: %w", err) + } + + // Binary protobuf decodes an empty value to a zero message without error, + // so validate the decoded identity against the key by reconstructing the + // key the record claims to live at. Reconstructing via actorDBKey keeps + // this correct for the atespace-scoped key shape and for ids containing + // ':'. Guards against an empty/corrupt key surfacing as a blank actor + // (GetActor performs the same check). + if actorDBKey(actor.GetAtespace(), actor.GetActorId()) != keys[i] { + return nil, fmt.Errorf("actor record at key %q does not match its key identity", keys[i]) } actors = append(actors, actor) } diff --git a/cmd/ateapi/internal/store/ateredis/ateredis_encoding_test.go b/cmd/ateapi/internal/store/ateredis/ateredis_encoding_test.go new file mode 100644 index 000000000..1bf04a80e --- /dev/null +++ b/cmd/ateapi/internal/store/ateredis/ateredis_encoding_test.go @@ -0,0 +1,168 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Tests for the binary-protobuf record encoding (issue #307). These cover the +// behaviour specific to the encoding switch: that records are stored as binary +// protobuf rather than protojson, and that the read paths reject empty/corrupt/ +// identity-mismatched values. The latter matters because, unlike protojson, +// proto.Unmarshal accepts an empty byte slice and yields a zero-valued message +// without error, so the key-identity checks are what catch a blank record. +package ateredis + +import ( + "testing" + + "google.golang.org/protobuf/proto" + + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" +) + +// TestCreateActor_StoresBinaryProtobuf locks in the binary-protobuf encoding so +// an accidental revert to protojson (which would still round-trip through the +// store API and pass the other tests) is caught: protojson values start with +// '{', binary protobuf ones do not. +func TestCreateActor_StoresBinaryProtobuf(t *testing.T) { + mr, s, ctx := setupTest(t) + defer mr.Close() + + actor := &ateapipb.Actor{ + ActorId: "session-1", + Atespace: "ns1", + ActorTemplateNamespace: "default", + ActorTemplateName: "test-template", + Status: ateapipb.Actor_STATUS_SUSPENDED, + } + if err := s.CreateActor(ctx, actor); err != nil { + t.Fatalf("CreateActor failed: %v", err) + } + + raw, err := s.rdb.Get(ctx, actorDBKey(actor.GetAtespace(), actor.GetActorId())).Bytes() + if err != nil { + t.Fatalf("raw Get failed: %v", err) + } + if len(raw) > 0 && raw[0] == '{' { + t.Errorf("stored actor value looks like JSON, expected binary protobuf: %q", raw) + } + if err := proto.Unmarshal(raw, &ateapipb.Actor{}); err != nil { + t.Errorf("stored actor value is not valid binary protobuf: %v", err) + } +} + +// TestCreateWorker_StoresBinaryProtobuf is the Worker counterpart. +func TestCreateWorker_StoresBinaryProtobuf(t *testing.T) { + mr, s, ctx := setupTest(t) + defer mr.Close() + + worker := &ateapipb.Worker{WorkerNamespace: "default", WorkerPool: "pool-1", WorkerPod: "pod-1"} + if err := s.CreateWorker(ctx, worker); err != nil { + t.Fatalf("CreateWorker failed: %v", err) + } + + raw, err := s.rdb.Get(ctx, workerDBKey("default", "pool-1", "pod-1")).Bytes() + if err != nil { + t.Fatalf("raw Get failed: %v", err) + } + if len(raw) > 0 && raw[0] == '{' { + t.Errorf("stored worker value looks like JSON, expected binary protobuf: %q", raw) + } + if err := proto.Unmarshal(raw, &ateapipb.Worker{}); err != nil { + t.Errorf("stored worker value is not valid binary protobuf: %v", err) + } +} + +func TestGetActor_RejectsEmptyValue(t *testing.T) { + mr, s, ctx := setupTest(t) + defer mr.Close() + + if err := s.rdb.Set(ctx, actorDBKey("ns1", "ghost"), "", 0).Err(); err != nil { + t.Fatalf("seeding empty value failed: %v", err) + } + if _, err := s.GetActor(ctx, "ns1", "ghost"); err == nil { + t.Errorf("expected GetActor to reject an empty value, got nil") + } +} + +func TestGetWorker_RejectsEmptyValue(t *testing.T) { + mr, s, ctx := setupTest(t) + defer mr.Close() + + if err := s.rdb.Set(ctx, workerDBKey("ns1", "pool1", "ghost"), "", 0).Err(); err != nil { + t.Fatalf("seeding empty value failed: %v", err) + } + if _, err := s.GetWorker(ctx, "ns1", "pool1", "ghost"); err == nil { + t.Errorf("expected GetWorker to reject an empty value, got nil") + } +} + +func TestListActors_RejectsEmptyValue(t *testing.T) { + mr, s, ctx := setupTest(t) + defer mr.Close() + + if err := s.rdb.Set(ctx, actorDBKey("ns1", "ghost"), "", 0).Err(); err != nil { + t.Fatalf("seeding empty value failed: %v", err) + } + if _, _, err := s.ListActors(ctx, "ns1", 1000, ""); err == nil { + t.Errorf("expected ListActors to reject an empty-value actor key, got nil") + } +} + +func TestListActors_RejectsIdentityMismatch(t *testing.T) { + mr, s, ctx := setupTest(t) + defer mr.Close() + + // A valid Actor proto, but stored under a key with a different actor id. + bytes, err := proto.Marshal(&ateapipb.Actor{ + ActorId: "real", Atespace: "ns1", ActorTemplateNamespace: "ns1", ActorTemplateName: "tmpl1", Version: 1, + }) + if err != nil { + t.Fatalf("marshal failed: %v", err) + } + if err := s.rdb.Set(ctx, actorDBKey("ns1", "wrong"), bytes, 0).Err(); err != nil { + t.Fatalf("seeding mismatched actor failed: %v", err) + } + if _, _, err := s.ListActors(ctx, "ns1", 1000, ""); err == nil { + t.Errorf("expected ListActors to reject an identity-mismatched actor, got nil") + } +} + +func TestListWorkers_RejectsEmptyValue(t *testing.T) { + mr, s, ctx := setupTest(t) + defer mr.Close() + + if err := s.rdb.Set(ctx, workerDBKey("ns1", "pool1", "ghost"), "", 0).Err(); err != nil { + t.Fatalf("seeding empty value failed: %v", err) + } + if _, err := s.ListWorkers(ctx); err == nil { + t.Errorf("expected ListWorkers to reject an empty-value worker key, got nil") + } +} + +func TestListWorkers_RejectsIdentityMismatch(t *testing.T) { + mr, s, ctx := setupTest(t) + defer mr.Close() + + // A valid Worker proto, but stored under a key with a different pod. + bytes, err := proto.Marshal(&ateapipb.Worker{ + WorkerNamespace: "ns1", WorkerPool: "pool1", WorkerPod: "real", Version: 1, + }) + if err != nil { + t.Fatalf("marshal failed: %v", err) + } + if err := s.rdb.Set(ctx, workerDBKey("ns1", "pool1", "wrong"), bytes, 0).Err(); err != nil { + t.Fatalf("seeding mismatched worker failed: %v", err) + } + if _, err := s.ListWorkers(ctx); err == nil { + t.Errorf("expected ListWorkers to reject an identity-mismatched worker, got nil") + } +}