adapter/sqs: validate peek cursor LastKey + add JSON tags to wire types

bootjp · bootjp · commit ddb4e6ec7b0f · 2026-05-21T22:28:58.000+09:00
Two r4 findings on PR #794: 1. Codex P2 (sqs_admin_peek.go:401). preparePeekCursor only validated Generation and partition indices; cursor.LastKey was trusted directly as the ScanAt start key, checked only against the upper bound. A forged LastKey below the queue's visibility-index prefix would let ScanAt walk unrelated key ranges, triggering up to Limit extra GetAt misses per call - read amplification against the leader. preparePeekCursor now rejects with ErrAdminSQSValidation when cursor.LastKey is not prefixed by sqsMsgVisPrefixForQueueDispatch(meta, queueName, Partition, Generation). The check catches three classes of forgery: completely arbitrary bytes (e.g. "aaaa"), a valid LastKey from a different queue (prefix encodes queue name), and a valid LastKey from a different partition of the same queue. preparePeekCursor's signature gains queueName; the only existing caller (AdminPeekQueue) and the existing unit tests were updated. 2. CodeRabbit Major (sqs_admin_peek.go:35-45, outside diff). AdminPeekedMessage lacked JSON tags, so json.Marshal emitted Go-style PascalCase ("MessageID", "BodyTruncated", ...) instead of the snake_case wire shape the design doc §3.5 specifies (the SPA's client adapter expects the spec'd form). Also empty Attributes / GroupID / DeduplicationID would serialize as "null" / "" instead of being omitted. JSON tags added with appropriate omitempty. Caller audit (semantic-change rule): preparePeekCursor signature changed from (cursor, meta, startPartition) to (cursor, meta, queueName, startPartition). One non-test caller (AdminPeekQueue at sqs_admin_peek.go:240) and 10 test call sites updated. The new queueName parameter is purely additive validation; it does not alter the cursor's stamping or generation-mismatch behavior. Tests: TestAdminPeekedMessage_JSONWireFormat / TestAdminPeekedMessage_JSONOmitsEmptyAttributes pin the JSON wire shape. TestPreparePeekCursor_PartitionOutOfRange gains three new sub-cases (forged LastKey, foreign-queue LastKey, mismatched- partition LastKey). The test function was rewritten as a table-driven loop so it stays under the cyclop budget.
diff --git a/adapter/sqs_admin_peek.go b/adapter/sqs_admin_peek.go
@@ -31,17 +31,21 @@ type AdminPeekedAttribute struct {
 	BinaryValue []byte `json:"binary_value,omitempty"`
 }
 
-// AdminPeekedMessage is one row in the peek result.
+// AdminPeekedMessage is one row in the peek result. JSON tags pin
+// the snake_case wire shape the design doc §3.5 specifies; without
+// them the encoder would emit Go-style PascalCase field names and
+// the SPA's client adapter would silently misparse every row.
+// CodeRabbit r4 caught the regression.
 type AdminPeekedMessage struct {
-	MessageID        string
-	Body             string                          // truncated per opts.BodyMaxBytes
-	BodyTruncated    bool                            // true when Body was cut
-	BodyOriginalSize int64                           // bytes in the original body, for display
-	SentTimestamp    time.Time                       // SQS SentTimestamp
-	ReceiveCount     int32                           // ApproximateReceiveCount
-	GroupID          string                          // FIFO MessageGroupId, empty for standard
-	DeduplicationID  string                          // FIFO MessageDeduplicationId, empty for standard
-	Attributes       map[string]AdminPeekedAttribute // typed SQS message attributes
+	MessageID        string                          `json:"message_id"`
+	Body             string                          `json:"body"`                       // truncated per opts.BodyMaxBytes
+	BodyTruncated    bool                            `json:"body_truncated"`             // true when Body was cut
+	BodyOriginalSize int64                           `json:"body_original_size"`         // bytes in the original body, for display
+	SentTimestamp    time.Time                       `json:"sent_timestamp"`             // SQS SentTimestamp
+	ReceiveCount     int32                           `json:"receive_count"`              // ApproximateReceiveCount
+	GroupID          string                          `json:"group_id,omitempty"`         // FIFO MessageGroupId, empty for standard
+	DeduplicationID  string                          `json:"deduplication_id,omitempty"` // FIFO MessageDeduplicationId, empty for standard
+	Attributes       map[string]AdminPeekedAttribute `json:"attributes,omitempty"`       // typed SQS message attributes
 }
 
 // AdminPeekMessageOptions controls a peek call. Zero values map to
@@ -233,7 +237,7 @@ func (s *SQSServer) AdminPeekQueue(
 	if !exists {
 		return nil, "", ErrAdminSQSNotFound
 	}
-	cursor, err = preparePeekCursor(cursor, meta, s.peekStartPartition(name, cursor, meta))
+	cursor, err = preparePeekCursor(cursor, meta, name, s.peekStartPartition(name, cursor, meta))
 	if err != nil {
 		return nil, "", err
 	}
@@ -304,7 +308,15 @@ func (s *SQSServer) peekStartPartition(queueName string, cursor *peekCursor, met
 // peekStartPartition wrapper around nextReceiveFanoutStart) so this
 // helper stays method-free for unit-testability. Non-partitioned and
 // perQueue-collapsed queues pass 0.
-func preparePeekCursor(cursor *peekCursor, meta *sqsQueueMeta, startPartition uint32) (*peekCursor, error) {
+//
+// queueName is consumed to validate cursor.LastKey on continuation
+// pages: a forged LastKey outside the queue's visibility-index
+// prefix would otherwise let an attacker start a ScanAt at any byte
+// offset in the leader's keyspace (Codex r4 P2). Continuation
+// cursors are admin-supplied and signed only by base64-encoding, so
+// every field must be validated against the live queue meta before
+// being used as a storage cursor.
+func preparePeekCursor(cursor *peekCursor, meta *sqsQueueMeta, queueName string, startPartition uint32) (*peekCursor, error) {
 	effective := effectivePartitionCount(meta)
 	if cursor == nil {
 		out := &peekCursor{
@@ -326,6 +338,13 @@ func preparePeekCursor(cursor *peekCursor, meta *sqsQueueMeta, startPartition ui
 			"admin peek: cursor partition index out of range (StartPartition=%d, Partition=%d, max=%d)",
 			cursor.StartPartition, cursor.Partition, effective)
 	}
+	if len(cursor.LastKey) > 0 {
+		expectedPrefix := sqsMsgVisPrefixForQueueDispatch(meta, queueName, cursor.Partition, meta.Generation)
+		if !bytes.HasPrefix(cursor.LastKey, expectedPrefix) {
+			return nil, errors.Wrap(ErrAdminSQSValidation,
+				"admin peek: cursor LastKey is outside the queue's visibility-index prefix")
+		}
+	}
 	return cursor, nil
 }
 
diff --git a/adapter/sqs_admin_peek_test.go b/adapter/sqs_admin_peek_test.go
@@ -727,6 +727,80 @@ func TestClampPeekBodyBytes(t *testing.T) {
 	}
 }
 
+// TestAdminPeekedMessage_JSONWireFormat pins the snake_case wire
+// shape the design doc §3.5 specifies. Without explicit JSON tags
+// the encoder emits Go-style PascalCase ("MessageID", "BodyTruncated"
+// …) and the SPA's client adapter silently misparses every row
+// (CodeRabbit r4 caught the regression).
+func TestAdminPeekedMessage_JSONWireFormat(t *testing.T) {
+	t.Parallel()
+	in := AdminPeekedMessage{
+		MessageID:        "m1",
+		Body:             "hello",
+		BodyTruncated:    true,
+		BodyOriginalSize: 42,
+		SentTimestamp:    time.UnixMilli(1_700_000_000_000).UTC(),
+		ReceiveCount:     3,
+		GroupID:          "g1",
+		DeduplicationID:  "d1",
+		Attributes: map[string]AdminPeekedAttribute{
+			"Source": {DataType: sqsAttributeBaseTypeString, StringValue: "checkout"},
+		},
+	}
+	raw, err := json.Marshal(in)
+	if err != nil {
+		t.Fatalf("marshal: %v", err)
+	}
+	var got map[string]any
+	if err := json.Unmarshal(raw, &got); err != nil {
+		t.Fatalf("unmarshal: %v", err)
+	}
+	wantKeys := []string{
+		"message_id", "body", "body_truncated", "body_original_size",
+		"sent_timestamp", "receive_count", "group_id", "deduplication_id",
+		"attributes",
+	}
+	for _, k := range wantKeys {
+		if _, ok := got[k]; !ok {
+			t.Fatalf("missing key %q in wire JSON; got=%v", k, got)
+		}
+	}
+	// PascalCase keys must NOT appear.
+	for _, k := range []string{"MessageID", "Body", "BodyTruncated", "GroupID", "Attributes"} {
+		if _, ok := got[k]; ok {
+			t.Fatalf("unwanted PascalCase key %q in wire JSON; got=%v", k, got)
+		}
+	}
+	// Nested AdminPeekedAttribute also uses snake_case.
+	attrs, _ := got["attributes"].(map[string]any)
+	src, _ := attrs["Source"].(map[string]any)
+	if _, ok := src["data_type"]; !ok {
+		t.Fatalf("Source attribute missing 'data_type' key; got=%v", src)
+	}
+}
+
+// TestAdminPeekedMessage_JSONOmitsEmptyAttributes pins the
+// omitempty contract: empty Attributes / GroupID / DeduplicationID
+// must NOT appear on the wire (the SPA renders the field's absence,
+// not a "null" or "{}" placeholder).
+func TestAdminPeekedMessage_JSONOmitsEmptyAttributes(t *testing.T) {
+	t.Parallel()
+	in := AdminPeekedMessage{MessageID: "m1", Body: "x"}
+	raw, err := json.Marshal(in)
+	if err != nil {
+		t.Fatalf("marshal: %v", err)
+	}
+	var got map[string]any
+	if err := json.Unmarshal(raw, &got); err != nil {
+		t.Fatalf("unmarshal: %v", err)
+	}
+	for _, k := range []string{"group_id", "deduplication_id", "attributes"} {
+		if _, ok := got[k]; ok {
+			t.Fatalf("empty %q must be omitted from wire JSON; got=%v", k, got)
+		}
+	}
+}
+
 // TestPreparePeekCursor_FreshCursor confirms the first-call cursor
 // stamps Generation and (for partitioned queues) StartPartition.
 func TestPreparePeekCursor_FreshCursor(t *testing.T) {
@@ -735,7 +809,7 @@ func TestPreparePeekCursor_FreshCursor(t *testing.T) {
 	t.Run("non-partitioned", func(t *testing.T) {
 		t.Parallel()
 		meta := &sqsQueueMeta{Generation: 7}
-		out, err := preparePeekCursor(nil, meta, 0)
+		out, err := preparePeekCursor(nil, meta, "test", 0)
 		if err != nil {
 			t.Fatalf("err: %v", err)
 		}
@@ -747,7 +821,7 @@ func TestPreparePeekCursor_FreshCursor(t *testing.T) {
 	t.Run("partitioned", func(t *testing.T) {
 		t.Parallel()
 		meta := &sqsQueueMeta{Generation: 11, PartitionCount: 4}
-		out, err := preparePeekCursor(nil, meta, 2)
+		out, err := preparePeekCursor(nil, meta, "test", 2)
 		if err != nil {
 			t.Fatalf("err: %v", err)
 		}
@@ -763,7 +837,7 @@ func TestPreparePeekCursor_GenerationMismatch(t *testing.T) {
 	t.Parallel()
 	stale := &peekCursor{V: peekCursorSchemaV1, Generation: 3}
 	meta := &sqsQueueMeta{Generation: 4}
-	_, err := preparePeekCursor(stale, meta, 0)
+	_, err := preparePeekCursor(stale, meta, "test", 0)
 	if !errors.Is(err, ErrAdminSQSValidation) {
 		t.Fatalf("want ErrAdminSQSValidation, got %v", err)
 	}
@@ -779,52 +853,50 @@ func TestPreparePeekCursor_GenerationMismatch(t *testing.T) {
 // endpoint).
 func TestPreparePeekCursor_PartitionOutOfRange(t *testing.T) {
 	t.Parallel()
-
-	t.Run("partitioned: StartPartition >= PartitionCount", func(t *testing.T) {
-		t.Parallel()
-		meta := &sqsQueueMeta{Generation: 1, PartitionCount: 4}
-		cursor := &peekCursor{V: peekCursorSchemaV1, Generation: 1, StartPartition: 99, Partition: 0}
-		_, err := preparePeekCursor(cursor, meta, 0)
-		if !errors.Is(err, ErrAdminSQSValidation) {
-			t.Fatalf("StartPartition=99/PC=4: want ErrAdminSQSValidation, got %v", err)
-		}
-	})
-
-	t.Run("partitioned: Partition >= PartitionCount", func(t *testing.T) {
-		t.Parallel()
-		meta := &sqsQueueMeta{Generation: 1, PartitionCount: 4}
-		cursor := &peekCursor{V: peekCursorSchemaV1, Generation: 1, StartPartition: 0, Partition: 99}
-		_, err := preparePeekCursor(cursor, meta, 0)
-		if !errors.Is(err, ErrAdminSQSValidation) {
-			t.Fatalf("Partition=99/PC=4: want ErrAdminSQSValidation, got %v", err)
-		}
-	})
-
-	t.Run("non-partitioned: non-zero StartPartition", func(t *testing.T) {
-		t.Parallel()
-		meta := &sqsQueueMeta{Generation: 1, PartitionCount: 0}
-		cursor := &peekCursor{V: peekCursorSchemaV1, Generation: 1, StartPartition: 1, Partition: 0}
-		_, err := preparePeekCursor(cursor, meta, 0)
-		if !errors.Is(err, ErrAdminSQSValidation) {
-			t.Fatalf("non-partitioned StartPartition=1: want ErrAdminSQSValidation, got %v", err)
-		}
-	})
-
-	t.Run("non-partitioned: non-zero Partition", func(t *testing.T) {
-		t.Parallel()
-		meta := &sqsQueueMeta{Generation: 1, PartitionCount: 1}
-		cursor := &peekCursor{V: peekCursorSchemaV1, Generation: 1, StartPartition: 0, Partition: 5}
-		_, err := preparePeekCursor(cursor, meta, 0)
-		if !errors.Is(err, ErrAdminSQSValidation) {
-			t.Fatalf("non-partitioned Partition=5: want ErrAdminSQSValidation, got %v", err)
-		}
-	})
+	pc4 := &sqsQueueMeta{Generation: 1, PartitionCount: 4}
+	pc0 := &sqsQueueMeta{Generation: 1, PartitionCount: 0}
+	pc1 := &sqsQueueMeta{Generation: 1, PartitionCount: 1}
+	// foreignKey / mismatchedKey share a "first byte not in the
+	// admin prefix" property with "aaaa" so the test exercises all
+	// three forged-LastKey classes the bounds check catches.
+	foreignKey := append(sqsMsgVisPrefixForQueueDispatch(pc4, "queue-X", 0, 1), 'X')
+	mismatchedKey := append(sqsMsgVisPrefixForQueueDispatch(pc4, "queue-A", 2, 1), 'X')
+
+	cases := []struct {
+		name   string
+		meta   *sqsQueueMeta
+		queue  string
+		cursor *peekCursor
+	}{
+		{"partitioned: StartPartition >= PartitionCount", pc4, "queue-A",
+			&peekCursor{V: peekCursorSchemaV1, Generation: 1, StartPartition: 99, Partition: 0}},
+		{"partitioned: Partition >= PartitionCount", pc4, "queue-A",
+			&peekCursor{V: peekCursorSchemaV1, Generation: 1, StartPartition: 0, Partition: 99}},
+		{"non-partitioned: non-zero StartPartition", pc0, "queue-A",
+			&peekCursor{V: peekCursorSchemaV1, Generation: 1, StartPartition: 1, Partition: 0}},
+		{"non-partitioned: non-zero Partition", pc1, "queue-A",
+			&peekCursor{V: peekCursorSchemaV1, Generation: 1, StartPartition: 0, Partition: 5}},
+		{"forged LastKey outside partition prefix (Codex r4 P2)", pc4, "queue-A",
+			&peekCursor{V: peekCursorSchemaV1, Generation: 1, StartPartition: 0, Partition: 0, LastKey: []byte("aaaa")}},
+		{"LastKey from a different queue", pc4, "queue-A",
+			&peekCursor{V: peekCursorSchemaV1, Generation: 1, StartPartition: 0, Partition: 0, LastKey: foreignKey}},
+		{"LastKey from a different partition", pc4, "queue-A",
+			&peekCursor{V: peekCursorSchemaV1, Generation: 1, StartPartition: 0, Partition: 0, LastKey: mismatchedKey}},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			_, err := preparePeekCursor(tc.cursor, tc.meta, tc.queue, 0)
+			if !errors.Is(err, ErrAdminSQSValidation) {
+				t.Fatalf("%s: want ErrAdminSQSValidation, got %v", tc.name, err)
+			}
+		})
+	}
 
 	t.Run("partitioned: in-range cursor accepted", func(t *testing.T) {
 		t.Parallel()
-		meta := &sqsQueueMeta{Generation: 1, PartitionCount: 4}
 		cursor := &peekCursor{V: peekCursorSchemaV1, Generation: 1, StartPartition: 3, Partition: 2}
-		out, err := preparePeekCursor(cursor, meta, 0)
+		out, err := preparePeekCursor(cursor, pc4, "queue-A", 0)
 		if err != nil {
 			t.Fatalf("in-range cursor: got err %v want nil", err)
 		}
@@ -958,7 +1030,7 @@ func TestPreparePeekCursor_PerQueueCollapse(t *testing.T) {
 	t.Run("fresh cursor on perQueue stamps StartPartition=0", func(t *testing.T) {
 		t.Parallel()
 		meta := &sqsQueueMeta{Generation: 1, PartitionCount: 4, FifoThroughputLimit: htfifoThroughputPerQueue}
-		out, err := preparePeekCursor(nil, meta, 2)
+		out, err := preparePeekCursor(nil, meta, "test", 2)
 		if err != nil {
 			t.Fatalf("err: %v", err)
 		}
@@ -971,7 +1043,7 @@ func TestPreparePeekCursor_PerQueueCollapse(t *testing.T) {
 		t.Parallel()
 		meta := &sqsQueueMeta{Generation: 1, PartitionCount: 4, FifoThroughputLimit: htfifoThroughputPerQueue}
 		cursor := &peekCursor{V: peekCursorSchemaV1, Generation: 1, StartPartition: 0, Partition: 2}
-		_, err := preparePeekCursor(cursor, meta, 0)
+		_, err := preparePeekCursor(cursor, meta, "test", 0)
 		if !errors.Is(err, ErrAdminSQSValidation) {
 			t.Fatalf("perQueue Partition=2 (raw PartitionCount=4 would allow): want ErrAdminSQSValidation, got %v", err)
 		}
