From af2560222ff1eadb4b97429e28847c12deff84cf Mon Sep 17 00:00:00 2001 From: Kaien Date: Sun, 9 Aug 2026 21:29:41 +0300 Subject: [PATCH] Harden Stage 5 subscription safety checks Fence stale ACK results by generation, make duplicate MarkProcessed harmless after confirmed ACK, and cover forged delivery, mutation, in-flight ACK, initial subscribe, and replacement edge cases. Co-authored-by: Cursor --- README.md | 16 +++ ack_tracker.go | 10 +- subscription.go | 5 +- subscription_test.go | 268 +++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 296 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index c50d8be..558f2b9 100644 --- a/README.md +++ b/README.md @@ -106,3 +106,19 @@ The SDK wraps implemented read/agent RPCs: Generic `PostMessage`, `CreateTopic`, and `UpdateTopic` are intentionally not exposed as ergonomic helpers because the current server returns `Unimplemented` for those fail-closed gRPC paths. + +## Lifecycle Notes + +`Client` owns and closes the gRPC connection it creates in `New`. If callers +provide an existing connection with `WithClientConn`, the SDK uses it but does +not close it; callers keep ownership of that connection. + +`Subscription.Run` is sequential: it calls the handler for one delivery at a +time and calls `MarkProcessed` only after the handler returns nil. Handler +errors stop `Run` and do not ACK that delivery. Automatic reconnect cancels the +subscription stream generation, but it cannot undo external handler side +effects; handlers should honor their context and remain idempotent. + +Calling `MarkProcessed` more than once for the same delivery is harmless. The +first call may advance the safe ACK boundary; later duplicate calls do not add +tracker entries, move the cursor backward, or ACK a caller-mutated sequence. diff --git a/ack_tracker.go b/ack_tracker.go index c888e55..9a6e655 100644 --- a/ack_tracker.go +++ b/ack_tracker.go @@ -13,10 +13,14 @@ type ackEntry struct { type ackTracker struct { entries []ackEntry index map[*deliveryToken]int + retired map[*deliveryToken]struct{} } func newACKTracker() ackTracker { - return ackTracker{index: make(map[*deliveryToken]int)} + return ackTracker{ + index: make(map[*deliveryToken]int), + retired: make(map[*deliveryToken]struct{}), + } } func (t *ackTracker) add(seq int64, token *deliveryToken) { @@ -27,6 +31,9 @@ func (t *ackTracker) add(seq int64, token *deliveryToken) { func (t *ackTracker) mark(token *deliveryToken) (int64, bool, bool) { i, ok := t.index[token] if !ok { + if _, retired := t.retired[token]; retired { + return t.safeBoundary(), true, true + } return 0, false, false } if t.entries[i].done { @@ -53,6 +60,7 @@ func (t *ackTracker) removeConfirmed(lastAck int64) { if !t.entries[n].done || t.entries[n].seq > lastAck { break } + t.retired[t.entries[n].token] = struct{}{} delete(t.index, t.entries[n].token) n++ } diff --git a/subscription.go b/subscription.go index 19d727b..b2ffe04 100644 --- a/subscription.go +++ b/subscription.go @@ -303,8 +303,9 @@ func (s *Subscription) runACK() { SubscriptionId: subID, }) if err != nil { - if s.ctx.Err() != nil || !isRetryableACKError(err) || !s.sameGeneration(gen) { - if s.ctx.Err() == nil && !isRetryableACKError(err) { + current := s.sameGeneration(gen) + if s.ctx.Err() != nil || !isRetryableACKError(err) || !current { + if s.ctx.Err() == nil && current && !isRetryableACKError(err) { s.setErr(err) s.cancel() } diff --git a/subscription_test.go b/subscription_test.go index 037ccd4..bb87a63 100644 --- a/subscription_test.go +++ b/subscription_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "io" "net" "strings" "sync" @@ -325,6 +326,101 @@ func TestOrderedCumulativeACKWithSequenceGaps(t *testing.T) { waitACKs(t, fake, 104, 111) } +func TestForgedDeliveryRejectedWithoutACK(t *testing.T) { + fake := newFake() + fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { + _ = stream.Send(&pb.DeliveryMessage{Ready: true}) + <-stream.Context().Done() + return stream.Context().Err() + } + c := newClientForTest(t, fake) + sub, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) + if err != nil { + t.Fatal(err) + } + err = sub.MarkProcessed(context.Background(), Delivery{ + Message: &pb.Message{Sequence: 999999}, + Replay: true, + }) + if !errors.Is(err, ErrUnknownDelivery) { + t.Fatalf("forged delivery MarkProcessed = %v", err) + } + fake.mu.Lock() + defer fake.mu.Unlock() + if len(fake.ackCalls) != 0 { + t.Fatalf("forged delivery produced ACKs: %+v", fake.ackCalls) + } +} + +func TestCrossSubscriptionDeliveryRejectedWithoutACK(t *testing.T) { + fake := newFake() + fake.subscribeFunc = func(req *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { + _ = stream.Send(&pb.DeliveryMessage{Ready: true}) + dm := msg(100) + dm.Message.TopicId = req.GetTopicId() + dm.SubscriptionId = "sub-" + req.GetTopicId() + _ = stream.Send(dm) + <-stream.Context().Done() + return stream.Context().Err() + } + c := newClientForTest(t, fake) + subA, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "a"}) + subB, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "b"}) + dA := recvDelivery(t, subA) + _ = recvDelivery(t, subB) + if err := subB.MarkProcessed(context.Background(), dA); !errors.Is(err, ErrUnknownDelivery) { + t.Fatalf("cross-subscription MarkProcessed = %v", err) + } + fake.mu.Lock() + defer fake.mu.Unlock() + if len(fake.ackCalls) != 0 { + t.Fatalf("cross-subscription delivery produced ACKs: %+v", fake.ackCalls) + } +} + +func TestMutatedDeliveryMessageCannotAlterACKSequence(t *testing.T) { + fake := newFake() + fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { + _ = stream.Send(&pb.DeliveryMessage{Ready: true}) + _ = stream.Send(msg(100)) + <-stream.Context().Done() + return stream.Context().Err() + } + c := newClientForTest(t, fake) + sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) + d := recvDelivery(t, sub) + d.Message.Sequence = 111 + if err := sub.MarkProcessed(context.Background(), d); err != nil { + t.Fatal(err) + } + waitACKs(t, fake, 100) +} + +func TestDuplicateMarkProcessedIsHarmless(t *testing.T) { + fake := newFake() + fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { + _ = stream.Send(&pb.DeliveryMessage{Ready: true}) + _ = stream.Send(msg(100)) + <-stream.Context().Done() + return stream.Context().Err() + } + c := newClientForTest(t, fake) + sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) + d := recvDelivery(t, sub) + if err := sub.MarkProcessed(context.Background(), d); err != nil { + t.Fatal(err) + } + waitACKs(t, fake, 100) + if err := sub.MarkProcessed(context.Background(), d); err != nil { + t.Fatalf("duplicate MarkProcessed = %v", err) + } + fake.mu.Lock() + defer fake.mu.Unlock() + if len(fake.ackCalls) != 1 || fake.ackCalls[0].seq != 100 { + t.Fatalf("duplicate MarkProcessed changed ACKs: %+v", fake.ackCalls) + } +} + func TestLargerOutOfOrderCompletionDoesNotSkipEarliestPending(t *testing.T) { fake := newFake() fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { @@ -447,6 +543,102 @@ func TestACKResponseLostAmbiguityReconnectUsesServerCursor(t *testing.T) { waitACKs(t, fake, 104, 104, 111) } +func TestACKPrefixGrowthWhileACKInFlightAdvancesAfterSuccess(t *testing.T) { + fake := newFake() + ack100Started := make(chan struct{}) + release100 := make(chan struct{}) + var once sync.Once + fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { + _ = stream.Send(&pb.DeliveryMessage{Ready: true}) + _ = stream.Send(msg(100)) + _ = stream.Send(msg(104)) + <-stream.Context().Done() + return stream.Context().Err() + } + fake.ackFunc = func(_ context.Context, req *pb.AcknowledgeRequest) (*pb.AcknowledgeResponse, error) { + fake.mu.Lock() + fake.ackCalls = append(fake.ackCalls, ackCall{topic: req.GetTopicId(), subID: req.GetSubscriptionId(), seq: req.GetLastAckSequence()}) + fake.mu.Unlock() + if req.GetLastAckSequence() == 100 { + once.Do(func() { close(ack100Started) }) + <-release100 + } + return &pb.AcknowledgeResponse{LastAckSequence: req.GetLastAckSequence()}, nil + } + c := newClientForTest(t, fake) + sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) + d100, d104 := recvDelivery(t, sub), recvDelivery(t, sub) + if err := sub.MarkProcessed(context.Background(), d100); err != nil { + t.Fatal(err) + } + <-ack100Started + if err := sub.MarkProcessed(context.Background(), d104); err != nil { + t.Fatal(err) + } + close(release100) + waitACKs(t, fake, 100, 104) +} + +func TestOldACKResultCannotCorruptReconnectedGeneration(t *testing.T) { + fake := newFake() + dropFirst := make(chan struct{}) + oldACKStarted := make(chan struct{}) + releaseOldACK := make(chan struct{}) + var oldOnce sync.Once + var calls int + fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { + calls++ + _ = stream.Send(&pb.DeliveryMessage{Ready: true}) + if calls == 1 { + _ = stream.Send(msg(100)) + <-dropFirst + return status.Error(codes.Unavailable, "generation A failed") + } + dm := msg(100) + dm.Replay = true + dm.SubscriptionId = "sub-2" + _ = stream.Send(dm) + <-stream.Context().Done() + return stream.Context().Err() + } + fake.ackFunc = func(_ context.Context, req *pb.AcknowledgeRequest) (*pb.AcknowledgeResponse, error) { + fake.mu.Lock() + fake.ackCalls = append(fake.ackCalls, ackCall{topic: req.GetTopicId(), subID: req.GetSubscriptionId(), seq: req.GetLastAckSequence()}) + fake.mu.Unlock() + if req.GetSubscriptionId() == "sub-1" { + oldOnce.Do(func() { close(oldACKStarted) }) + <-releaseOldACK + return nil, status.Error(codes.Aborted, "old subscription_id rejected") + } + return &pb.AcknowledgeResponse{LastAckSequence: req.GetLastAckSequence()}, nil + } + c := newClientForTest(t, fake, + WithReconnectBackoff(10*time.Millisecond, 10*time.Millisecond), + withSleeperForTest(sleepContext), + ) + sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) + old100 := recvDelivery(t, sub) + if err := sub.MarkProcessed(context.Background(), old100); err != nil { + t.Fatal(err) + } + <-oldACKStarted + close(dropFirst) + new100 := recvDelivery(t, sub) + close(releaseOldACK) + if err := sub.MarkProcessed(context.Background(), new100); err != nil { + t.Fatal(err) + } + waitACKs(t, fake, 100, 100) + if err := sub.Err(); err != nil { + t.Fatalf("old ACK result corrupted generation B: %v", err) + } + fake.mu.Lock() + defer fake.mu.Unlock() + if fake.ackCalls[0].subID != "sub-1" || fake.ackCalls[1].subID != "sub-2" { + t.Fatalf("ACKs used wrong subscription ids: %+v", fake.ackCalls) + } +} + func TestReconnectReplayAndStaleGenerationFencing(t *testing.T) { fake := newFake() var calls int @@ -507,6 +699,48 @@ func TestStreamReplacementAbortedIsTerminalNoReconnect(t *testing.T) { } } +func TestReplacementBetweenTwoSDKSubscriptionsDoesNotPingPong(t *testing.T) { + fake := newFake() + replaceFirst := make(chan struct{}) + secondReady := make(chan struct{}) + fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { + fake.mu.Lock() + call := fake.subscribeCount + fake.mu.Unlock() + _ = stream.Send(&pb.DeliveryMessage{Ready: true}) + if call == 1 { + <-replaceFirst + return status.Error(codes.Aborted, "subscription replaced") + } + close(secondReady) + <-stream.Context().Done() + return stream.Context().Err() + } + c := newClientForTest(t, fake) + first, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) + if err != nil { + t.Fatal(err) + } + second, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) + if err != nil { + t.Fatal(err) + } + <-secondReady + close(replaceFirst) + waitDone(t, first) + if !errors.Is(first.Err(), ErrSubscriptionReplaced) { + t.Fatalf("first Err = %v", first.Err()) + } + if err := second.Close(); err != nil { + t.Fatal(err) + } + fake.mu.Lock() + defer fake.mu.Unlock() + if fake.subscribeCount != 2 { + t.Fatalf("replacement caused reconnect ping-pong: %d subscribes", fake.subscribeCount) + } +} + func TestAuthTerminalErrorsDoNotEndlesslyReconnect(t *testing.T) { for _, code := range []codes.Code{codes.Unauthenticated, codes.PermissionDenied} { t.Run(code.String(), func(t *testing.T) { @@ -527,6 +761,40 @@ func TestAuthTerminalErrorsDoNotEndlesslyReconnect(t *testing.T) { } } +func TestInitialSubscribeFailuresBeforeReadyReturnImmediately(t *testing.T) { + cases := []struct { + name string + err error + code codes.Code + }{ + {"Unavailable", status.Error(codes.Unavailable, "transient outage"), codes.Unavailable}, + {"EOF", io.EOF, codes.Unknown}, + {"DeadlineExceeded", status.Error(codes.DeadlineExceeded, "deadline"), codes.DeadlineExceeded}, + {"Unauthenticated", status.Error(codes.Unauthenticated, "bad token"), codes.Unauthenticated}, + {"PermissionDenied", status.Error(codes.PermissionDenied, "denied"), codes.PermissionDenied}, + {"NotFound", status.Error(codes.NotFound, "missing topic"), codes.NotFound}, + {"InvalidArgument", status.Error(codes.InvalidArgument, "bad topic"), codes.InvalidArgument}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + fake := newFake() + fake.subscribeFunc = func(_ *pb.SubscribeRequest, _ pb.ChatService_SubscribeServer) error { + return tc.err + } + c := newClientForTest(t, fake) + _, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) + if status.Code(err) != tc.code { + t.Fatalf("Subscribe error = %v, code %v want %v", err, status.Code(err), tc.code) + } + fake.mu.Lock() + defer fake.mu.Unlock() + if fake.subscribeCount != 1 { + t.Fatalf("initial failure retried before READY: %d subscribes", fake.subscribeCount) + } + }) + } +} + func TestTransientReconnect(t *testing.T) { fake := newFake() var calls int