Merge pull request #1267 from nicktrav/nickt.compaction-invalidating-…

…iter

db: improve use-after-free test coverage for compactions

compaction_iter.go

@@ -122,6 +122,27 @@ import (
 // to take the range tombstones into consideration when outputting normal
 // keys. Just as with point deletions, a range deletion covering an entry can
 // cause the entry to be elided.
+//
+// A note on the stability of keys and values.
+//
+// The stability guarantees of keys and values returned by the iterator tree
+// that backs a compactionIter is nuanced and care must be taken when
+// referencing any returned items.
+//
+// Keys and values returned by exported functions (i.e. First, Next, etc.) have
+// lifetimes that fall into two categories:
+//
+// Lifetime valid for duration of compaction. Range deletion keys and values are
+// stable for the duration of the compaction, due to way in which a
+// compactionIter is typically constructed (i.e. via (*compaction).newInputIter,
+// which wraps the iterator over the range deletion block in a noCloseIter,
+// preventing the release of the backing memory until the compaction is
+// finished).
+//
+// Lifetime limited to duration of sstable block liveness. Point keys (SET, DEL,
+// etc.) and values must be cloned / copied following the return from the
+// exported function, and before a subsequent call to Next advances the iterator
+// and mutates the contents of the returned key and value.
 type compactionIter struct {
 	cmp   Compare
 	merge Merge
@@ -265,6 +286,20 @@ func (i *compactionIter) Next() (*InternalKey, []byte) {
 			// Although, note that `skip` may already be true before reaching here
 			// due to an earlier key in the stripe. Then it is fine to leave it set
 			// to true, as the earlier key must have had a higher sequence number.
+			//
+			// NOTE: there is a subtle invariant violation here in that calling
+			// saveKey and returning a reference to the temporary slice violates
+			// the stability guarantee for range deletion keys. A potential
+			// mediation could return the original iterKey and iterValue
+			// directly, as the backing memory is guaranteed to be stable until
+			// the compaction completes. The violation here is only minor in
+			// that the caller immediately clones the range deletion InternalKey
+			// when passing the key to the deletion fragmenter (see the
+			// call-site in compaction.go).
+			// TODO(travers): address this violation by removing the call to
+			// saveKey and instead return the original iterKey and iterValue.
+			// This goes against the comment on i.key in the struct, and
+			// therefore warrants some investigation.
 			i.saveKey()
 			i.value = i.iterValue
 			i.valid = true
@@ -412,6 +447,12 @@ const (
 
 // nextInStripe advances the iterator and returns one of the above const ints
 // indicating how its state changed.
+//
+// Calls to nextInStripe must be preceded by a call to saveKey to retain a
+// temporary reference to the original key, so that forward iteration can
+// proceed with a reference to the original key. Care should be taken to avoid
+// overwriting or mutating the saved key or value before they have been returned
+// to the caller of the exported function (i.e. the caller of Next, First, etc.)
 func (i *compactionIter) nextInStripe() stripeChangeType {
 	if !i.iterNext() {
 		return newStripe

compaction_iter_test.go

@@ -80,6 +80,13 @@ func TestCompactionIter(t *testing.T) {
 	var allowZeroSeqnum bool
 
 	newIter := func() *compactionIter {
+		// To adhere to the existing assumption that range deletion blocks in
+		// SSTables are not released while iterating, and therefore not
+		// susceptible to use-after-free bugs, we skip the zeroing of
+		// RangeDelete keys.
+		iter := newInvalidatingIter(&fakeIter{keys: keys, vals: vals})
+		iter.ignoreKind(InternalKeyKindRangeDelete)
+
 		return newCompactionIter(
 			DefaultComparer.Compare,
 			DefaultComparer.FormatKey,
@@ -88,7 +95,7 @@ func TestCompactionIter(t *testing.T) {
 				m.buf = append(m.buf, value...)
 				return m, nil
 			},
-			&fakeIter{keys: keys, vals: vals},
+			iter,
 			snapshots,
 			&rangedel.Fragmenter{},
 			allowZeroSeqnum,

internal/cache/value_invariants.go

@@ -43,9 +43,9 @@ func newValue(n int) *Value {
 func (v *Value) free() {
 	// When "invariants" are enabled set the value contents to 0xff in order to
 	// cache use-after-free bugs.
-	// for i := range v.buf {
-	// 	v.buf[i] = 0xff
-	// }
+	for i := range v.buf {
+		v.buf[i] = 0xff
+	}
 	manual.Free(v.buf)
 	// Setting Value.buf to nil is needed for correctness of the leak checking
 	// that is performed when the "invariants" or "tracing" build tags are

iterator_test.go

@@ -215,25 +215,22 @@ func (f *fakeIter) SetBounds(lower, upper []byte) {
 // invalidatingIter tests unsafe key/value slice reuse by modifying the last
 // returned key/value to all 1s.
 type invalidatingIter struct {
-	iter      internalIterator
-	lastKey   *InternalKey
-	lastValue []byte
+	iter        internalIterator
+	lastKey     *InternalKey
+	lastValue   []byte
+	ignoreKinds [base.InternalKeyKindMax]bool
 }
 
 func newInvalidatingIter(iter internalIterator) *invalidatingIter {
 	return &invalidatingIter{iter: iter}
 }
 
+func (i *invalidatingIter) ignoreKind(kind base.InternalKeyKind) {
+	i.ignoreKinds[kind] = true
+}
+
 func (i *invalidatingIter) update(key *InternalKey, value []byte) (*InternalKey, []byte) {
-	if i.lastKey != nil {
-		for j := range i.lastKey.UserKey {
-			i.lastKey.UserKey[j] = 0xff
-		}
-		i.lastKey.Trailer = 0
-	}
-	for j := range i.lastValue {
-		i.lastValue[j] = 0xff
-	}
+	i.zeroLast()
 
 	if key == nil {
 		i.lastKey = nil
@@ -248,6 +245,25 @@ func (i *invalidatingIter) update(key *InternalKey, value []byte) (*InternalKey,
 	return i.lastKey, i.lastValue
 }
 
+func (i *invalidatingIter) zeroLast() {
+	if i.lastKey == nil {
+		return
+	}
+	if i.ignoreKinds[i.lastKey.Kind()] {
+		return
+	}
+
+	if i.lastKey != nil {
+		for j := range i.lastKey.UserKey {
+			i.lastKey.UserKey[j] = 0xff
+		}
+		i.lastKey.Trailer = 0
+	}
+	for j := range i.lastValue {
+		i.lastValue[j] = 0xff
+	}
+}
+
 func (i *invalidatingIter) SeekGE(key []byte) (*InternalKey, []byte) {
 	return i.update(i.iter.SeekGE(key))
 }

sstable/block.go

@@ -187,6 +187,10 @@ type blockEntry struct {
 // tombstones are always encoded with a restart interval of 1. This per-block
 // key stability guarantee is sufficient for range tombstones as they are
 // always encoded in a single block.
+//
+// A blockIter also provides a value stability guarantee for range deletions
+// since there is only a single range deletion block per sstable and the
+// blockIter will not release the bytes for the block until it is closed.
 type blockIter struct {
 	cmp Compare
 	// offset is the byte index that marks where the current key/value is