storage/pebbleiter: mangle unsafe buffers between positioning methods

Randomly clobber the key and value buffers in order to ensure lifetimes are
respected. This commit extends the `pebbleiter.assertionIter` type used in
`crdb_test` builds that sits between pkg/storage and the pebble.Iterator type.
It now will randomly zero the buffers used to return the previous iterator
position's key and value, to ensure code isn't improperly relying on the
stability of the these keys.

Epic: None
Close #86657.
Release note: None

pkg/storage/pebbleiter/BUILD.bazel

@@ -13,7 +13,10 @@ go_library(
     }),
     importpath = "github.com/cockroachdb/cockroach/pkg/storage/pebbleiter",
     visibility = ["//visibility:public"],
-    deps = ["@com_github_cockroachdb_pebble//:pebble"],
+    deps = [
+        "@com_github_cockroachdb_errors//:errors",
+        "@com_github_cockroachdb_pebble//:pebble",
+    ],
 )
 
 REMOVE_GO_BUILD_CONSTRAINTS = "cat $< | grep -v '//go:build' | grep -v '// +build' > $@"

pkg/storage/pebbleiter/crdb_test_on.go

@@ -13,7 +13,12 @@
 
 package pebbleiter
 
-import "github.com/cockroachdb/pebble"
+import (
+	"math/rand"
+
+	"github.com/cockroachdb/errors"
+	"github.com/cockroachdb/pebble"
+)
 
 // Iterator wraps the *pebble.Iterator in crdb_test builds with an assertionIter
 // that detects when Close is called on the iterator twice.  Double closes are
@@ -31,6 +36,23 @@ func MaybeWrap(iter *pebble.Iterator) Iterator {
 type assertionIter struct {
 	*pebble.Iterator
 	closed bool
+	// unsafeBufs hold buffers used for returning values with short lifetimes to
+	// the caller. To assert that the client is respecting the lifetimes,
+	// assertionIter mangles the buffers as soon as the associated lifetime
+	// expires. This is the same technique applied by the unsafeMVCCIterator in
+	// pkg/storage, but this time applied at the API boundary between
+	// pkg/storage and Pebble.
+	//
+	// unsafeBufs holds two buffers per-key type and an index indicating which
+	// are currently in use. This is used to randomly switch to a different
+	// buffer, ensuring that the buffer(s) returned to the caller for the
+	// previous iterator position are garbage (as opposed to just state
+	// corresponding to the current iterator position).
+	unsafeBufs struct {
+		idx int
+		key [2][]byte
+		val [2][]byte
+	}
 }
 
 func (i *assertionIter) Clone(cloneOpts pebble.CloneOptions) (Iterator, error) {
@@ -43,8 +65,111 @@ func (i *assertionIter) Clone(cloneOpts pebble.CloneOptions) (Iterator, error) {
 
 func (i *assertionIter) Close() error {
 	if i.closed {
-		panic("pebble.Iterator already closed")
+		panic(errors.AssertionFailedf("pebble.Iterator already closed"))
 	}
 	i.closed = true
 	return i.Iterator.Close()
 }
+
+func (i *assertionIter) Key() []byte {
+	if !i.Valid() {
+		panic(errors.AssertionFailedf("Key() called on !Valid() pebble.Iterator"))
+	}
+	idx := i.unsafeBufs.idx
+	i.unsafeBufs.key[idx] = append(i.unsafeBufs.key[idx][:0], i.Iterator.Key()...)
+	return i.unsafeBufs.key[idx]
+}
+
+func (i *assertionIter) Value() []byte {
+	if !i.Valid() {
+		panic(errors.AssertionFailedf("Value() called on !Valid() pebble.Iterator"))
+	}
+	idx := i.unsafeBufs.idx
+	i.unsafeBufs.val[idx] = append(i.unsafeBufs.val[idx][:0], i.Iterator.Value()...)
+	return i.unsafeBufs.val[idx]
+}
+
+func (i *assertionIter) LazyValue() pebble.LazyValue {
+	if !i.Valid() {
+		panic(errors.AssertionFailedf("LazyValue() called on !Valid() pebble.Iterator"))
+	}
+	return i.Iterator.LazyValue()
+}
+
+func (i *assertionIter) First() bool {
+	i.maybeMangleBufs()
+	return i.Iterator.First()
+}
+
+func (i *assertionIter) SeekGE(key []byte) bool {
+	i.maybeMangleBufs()
+	return i.Iterator.SeekGE(key)
+}
+
+func (i *assertionIter) SeekGEWithLimit(key []byte, limit []byte) pebble.IterValidityState {
+	i.maybeMangleBufs()
+	return i.Iterator.SeekGEWithLimit(key, limit)
+}
+
+func (i *assertionIter) SeekPrefixGE(key []byte) bool {
+	i.maybeMangleBufs()
+	return i.Iterator.SeekPrefixGE(key)
+}
+
+func (i *assertionIter) Next() bool {
+	i.maybeMangleBufs()
+	return i.Iterator.Next()
+}
+
+func (i *assertionIter) NextWithLimit(limit []byte) pebble.IterValidityState {
+	i.maybeMangleBufs()
+	return i.Iterator.NextWithLimit(limit)
+}
+
+func (i *assertionIter) NextPrefix() bool {
+	i.maybeMangleBufs()
+	return i.Iterator.NextPrefix()
+}
+
+func (i *assertionIter) Last() bool {
+	i.maybeMangleBufs()
+	return i.Iterator.Last()
+}
+
+func (i *assertionIter) SeekLT(key []byte) bool {
+	i.maybeMangleBufs()
+	return i.Iterator.SeekLT(key)
+}
+
+func (i *assertionIter) SeekLTWithLimit(key []byte, limit []byte) pebble.IterValidityState {
+	i.maybeMangleBufs()
+	return i.Iterator.SeekLTWithLimit(key, limit)
+}
+
+func (i *assertionIter) Prev() bool {
+	i.maybeMangleBufs()
+	return i.Iterator.Prev()
+}
+
+func (i *assertionIter) PrevWithLimit(limit []byte) pebble.IterValidityState {
+	i.maybeMangleBufs()
+	return i.Iterator.PrevWithLimit(limit)
+}
+
+// maybeMangleBufs trashes the contents of buffers used to return unsafe values
+// to the caller. This is used to ensure that the client respects the Pebble
+// iterator interface and the lifetimes of buffers it returns.
+func (i *assertionIter) maybeMangleBufs() {
+	if rand.Intn(2) == 0 {
+		idx := i.unsafeBufs.idx
+		for _, b := range [...][]byte{i.unsafeBufs.key[idx], i.unsafeBufs.val[idx]} {
+			for i := range b {
+				b[i] = 0
+			}
+		}
+		if rand.Intn(2) == 0 {
+			// Switch to a new buffer for the next iterator position.
+			i.unsafeBufs.idx = (i.unsafeBufs.idx + 1) % 2
+		}
+	}
+}