Skip to content

Commit

Permalink
storage/pebbleiter: mangle unsafe buffers between positioning methods
Browse files Browse the repository at this point in the history
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 cockroachdb#86657.
Release note: None
  • Loading branch information
jbowens committed Feb 8, 2023
1 parent b031933 commit 88392b2
Show file tree
Hide file tree
Showing 2 changed files with 131 additions and 3 deletions.
5 changes: 4 additions & 1 deletion pkg/storage/pebbleiter/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -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' > $@"
Expand Down
129 changes: 127 additions & 2 deletions pkg/storage/pebbleiter/crdb_test_on.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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) {
Expand All @@ -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("Value() called on !Valid() pebble.Iterator"))
}
return i.Iterator.LazyValue()
}

func (i *assertionIter) First() bool {
i.mangleBufs()
return i.Iterator.First()
}

func (i *assertionIter) SeekGE(key []byte) bool {
i.mangleBufs()
return i.Iterator.SeekGE(key)
}

func (i *assertionIter) SeekGEWithLimit(key []byte, limit []byte) pebble.IterValidityState {
i.mangleBufs()
return i.Iterator.SeekGEWithLimit(key, limit)
}

func (i *assertionIter) SeekPrefixGE(key []byte) bool {
i.mangleBufs()
return i.Iterator.SeekPrefixGE(key)
}

func (i *assertionIter) Next() bool {
i.mangleBufs()
return i.Iterator.Next()
}

func (i *assertionIter) NextWithLimit(limit []byte) pebble.IterValidityState {
i.mangleBufs()
return i.Iterator.NextWithLimit(limit)
}

func (i *assertionIter) NextPrefix() bool {
i.mangleBufs()
return i.Iterator.NextPrefix()
}

func (i *assertionIter) Last() bool {
i.mangleBufs()
return i.Iterator.Last()
}

func (i *assertionIter) SeekLT(key []byte) bool {
i.mangleBufs()
return i.Iterator.SeekLT(key)
}

func (i *assertionIter) SeekLTWithLimit(key []byte, limit []byte) pebble.IterValidityState {
i.mangleBufs()
return i.Iterator.SeekLTWithLimit(key, limit)
}

func (i *assertionIter) Prev() bool {
i.mangleBufs()
return i.Iterator.Prev()
}

func (i *assertionIter) PrevWithLimit(limit []byte) pebble.IterValidityState {
i.mangleBufs()
return i.Iterator.PrevWithLimit(limit)
}

// mangleBufs 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) mangleBufs() {
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
}
}
}

0 comments on commit 88392b2

Please sign in to comment.