colblk: add more bound checks

Add invariants-only bound checking, especially before accessing
`UnsafeOffset`s. We also change to using `At2()` instead of two `At()`
calls in a few places.

Author: RaduBerinde

internal/invariants/off.go

@@ -49,3 +49,7 @@ func (*Value[V]) Get() V {
 
 // Set the value; no-op in non-invariant builds.
 func (*Value[V]) Set(v V) {}
+
+// CheckBounds panics if the index is not in the range [0, n). No-op in
+// non-invariant builds.
+func CheckBounds(i int, n int) {}

internal/invariants/on.go

@@ -6,7 +6,10 @@
 
 package invariants
 
-import "math/rand/v2"
+import (
+	"fmt"
+	"math/rand/v2"
+)
 
 // Sometimes returns true percent% of the time if invariants are Enabled (i.e.
 // we were built with the "invariants" or "race" build tags). Otherwise, always
@@ -72,3 +75,11 @@ func (v *Value[V]) Get() V {
 func (v *Value[V]) Set(inner V) {
 	v.v = inner
 }
+
+// CheckBounds panics if the index is not in the range [0, n). No-op in
+// non-invariant builds.
+func CheckBounds(i int, n int) {
+	if i < 0 || i >= n {
+		panic(fmt.Sprintf("index %d out of bounds [0, %d)", i, n))
+	}
+}

sstable/colblk/bitmap.go

@@ -69,6 +69,7 @@ func (b Bitmap) At(i int) bool {
 		// zero bitmap case.
 		return false
 	}
+	invariants.CheckBounds(i, b.bitCount)
 	val := b.data.At(int(uint(i) >> 6)) // aka At(i/64)
 	return val&(1<<(uint(i)&63)) != 0
 }
@@ -96,6 +97,10 @@ func (b Bitmap) SeekSetBitGE(i int) int {
 	// on the summary word to get the index of which word has a set bit, if any.
 	summaryTableOffset, summaryTableEnd := b.summaryTableBounds()
 	summaryWordIdx := summaryTableOffset + wordIdx>>6
+	if invariants.Enabled {
+		sz := bitmapRequiredSize(b.bitCount)
+		invariants.CheckBounds(summaryTableEnd-1, sz)
+	}
 	summaryNext := nextBitInWord(b.data.At(summaryWordIdx), uint(wordIdx%64)+1)
 	// If [summaryNext] == 64, then there are no set bits in any of the earlier
 	// words represented by the summary word at [summaryWordIdx]. In that case,
@@ -171,11 +176,11 @@ func (b Bitmap) SeekSetBitLE(i int) int {
 // in the bitmap. The i parameter must be in [0, bitCount). Returns the number
 // of bits represented by the bitmap if no next bit is unset.
 func (b Bitmap) SeekUnsetBitGE(i int) int {
+	invariants.CheckBounds(i, b.bitCount)
 	if b.data.ptr == nil {
 		// Zero bitmap case.
 		return i
 	}
-
 	wordIdx := i >> 6 // i/64
 	// If the there's a bit ≥ i unset in the same word, return it.
 	if next := nextBitInWord(^b.data.At(wordIdx), uint(i)&63); next < 64 {
@@ -199,6 +204,7 @@ func (b Bitmap) SeekUnsetBitGE(i int) int {
 // bitmap. The i parameter must be in [0, bitCount). Returns -1 if no previous
 // bit is unset.
 func (b Bitmap) SeekUnsetBitLE(i int) int {
+	invariants.CheckBounds(i, b.bitCount)
 	if b.data.ptr == nil {
 		// Zero bitmap case.
 		return i

sstable/colblk/data_block.go

@@ -1336,6 +1336,7 @@ func (i *DataBlockIter) Next() *base.InternalKV {
 			i.kv.K.SetSeqNum(base.SeqNum(n))
 		}
 	}
+	invariants.CheckBounds(i.row, i.d.values.slices)
 	// Inline i.d.values.At(row).
 	v := i.d.values.slice(i.d.values.offsets.At2(i.row))
 	if i.d.isValueExternal.At(i.row) {
@@ -1504,9 +1505,10 @@ func (i *DataBlockIter) decodeRow() *base.InternalKV {
 				i.kv.K.SetSeqNum(base.SeqNum(n))
 			}
 		}
+		invariants.CheckBounds(i.row, i.d.values.slices)
 		// Inline i.d.values.At(row).
-		startOffset := i.d.values.offsets.At(i.row)
-		v := unsafe.Slice((*byte)(i.d.values.ptr(startOffset)), i.d.values.offsets.At(i.row+1)-startOffset)
+		v := i.d.values.slice(i.d.values.offsets.At2(i.row))
+		invariants.CheckBounds(i.row, i.d.values.slices)
 		if i.d.isValueExternal.At(i.row) {
 			i.kv.V = i.getLazyValuer.GetInternalValueForPrefixAndValueHandle(v)
 		} else {

sstable/colblk/prefix_bytes.go

@@ -282,6 +282,7 @@ func (i *PrefixBytesIter) Init(maxKeyLength int, syntheticPrefix block.Synthetic
 func (b *PrefixBytes) SetAt(it *PrefixBytesIter, i int) {
 	// Determine the offset and length of the bundle prefix.
 	bundleOffsetIndex := b.bundleOffsetIndexForRow(i)
+	invariants.CheckBounds(bundleOffsetIndex, b.rawBytes.slices)
 	bundleOffsetStart, bundleOffsetEnd := b.rawBytes.offsets.At2(bundleOffsetIndex)
 	bundlePrefixLen := bundleOffsetEnd - bundleOffsetStart
 
@@ -329,6 +330,7 @@ func (b *PrefixBytes) SetNext(it *PrefixBytesIter) {
 	// If the next row is in the same bundle, we can take a fast path of only
 	// updating the per-row suffix.
 	if it.offsetIndex < it.nextBundleOffsetIndex {
+		invariants.CheckBounds(it.offsetIndex, b.rawBytes.slices)
 		rowSuffixStart, rowSuffixEnd := b.rawBytes.offsets.At2(it.offsetIndex)
 		rowSuffixLen := rowSuffixEnd - rowSuffixStart
 		if rowSuffixLen == 0 {
@@ -351,6 +353,7 @@ func (b *PrefixBytes) SetNext(it *PrefixBytesIter) {
 	// The offsetIndex is currently pointing to the start of the new bundle
 	// prefix. Increment it to point at the start of the new row suffix.
 	it.offsetIndex++
+	invariants.CheckBounds(it.offsetIndex, b.rawBytes.slices)
 	rowSuffixStart, rowSuffixEnd := b.rawBytes.offsets.At2(it.offsetIndex)
 	rowSuffixLen := rowSuffixEnd - rowSuffixStart
 
@@ -391,15 +394,17 @@ func (b *PrefixBytes) SharedPrefix() []byte {
 // prefix for the column. The returned slice should not be mutated.
 func (b *PrefixBytes) RowBundlePrefix(row int) []byte {
 	i := b.bundleOffsetIndexForRow(row)
-	return b.rawBytes.slice(b.rawBytes.offsets.At(i), b.rawBytes.offsets.At(i+1))
+	invariants.CheckBounds(i, b.rawBytes.slices)
+	return b.rawBytes.slice(b.rawBytes.offsets.At2(i))
 }
 
 // BundlePrefix returns the prefix of the i-th bundle in the column. The
 // provided i must be in the range [0, BundleCount()). The returned slice should
 // not be mutated.
 func (b *PrefixBytes) BundlePrefix(i int) []byte {
 	j := b.offsetIndexByBundleIndex(i)
-	return b.rawBytes.slice(b.rawBytes.offsets.At(j), b.rawBytes.offsets.At(j+1))
+	invariants.CheckBounds(j, b.rawBytes.slices)
+	return b.rawBytes.slice(b.rawBytes.offsets.At2(j))
 }
 
 // RowSuffix returns a []byte of the suffix unique to the row. A row's full key
@@ -415,6 +420,7 @@ func (b *PrefixBytes) RowSuffix(row int) []byte {
 // accounting for duplicate keys. It takes the index of the row, and the value
 // of rowSuffixIndex(row).
 func (b *PrefixBytes) rowSuffixOffsets(row, i int) (low uint32, high uint32) {
+	invariants.CheckBounds(i, b.rawBytes.slices)
 	// Retrieve the low and high offsets indicating the start and end of the
 	// row's suffix slice.
 	low, high = b.rawBytes.offsets.At2(i)
@@ -493,6 +499,7 @@ func (b *PrefixBytes) Search(k []byte) (rowIndex int, isEqual bool) {
 		//     offset(j)             offset(j+1)                       offset(j+2)
 		//
 		j := b.offsetIndexByBundleIndex(h)
+		invariants.CheckBounds(j+1, b.rawBytes.slices)
 		bundleFirstKey := b.rawBytes.slice(b.rawBytes.offsets.At(j), b.rawBytes.offsets.At(j+2))
 		c = bytes.Compare(k, bundleFirstKey)
 		switch {
@@ -516,7 +523,8 @@ func (b *PrefixBytes) Search(k []byte) (rowIndex int, isEqual bool) {
 	// among them, but if the seek key doesn't share the previous bundle's
 	// prefix there's no need.
 	j := b.offsetIndexByBundleIndex(bi - 1)
-	bundlePrefix := b.rawBytes.slice(b.rawBytes.offsets.At(j), b.rawBytes.offsets.At(j+1))
+	invariants.CheckBounds(j, b.rawBytes.slices)
+	bundlePrefix := b.rawBytes.slice(b.rawBytes.offsets.At2(j))
 
 	// The row we are looking for might still be in the previous bundle even
 	// though the seek key is greater than the first key. This is possible only
@@ -553,6 +561,7 @@ func (b *PrefixBytes) Search(k []byte) (rowIndex int, isEqual bool) {
 		// The beginning of the zero-indexed i-th key of the bundle is at
 		// offset(j+i+1).
 		//
+		invariants.CheckBounds(j+h+1, b.rawBytes.slices)
 		hStart, hEnd := b.rawBytes.offsets.At2(j + h + 1)
 		// There's a complication with duplicate keys. When keys are repeated,
 		// the PrefixBytes encoding avoids re-encoding the duplicate key,

sstable/colblk/raw_bytes.go

@@ -11,6 +11,7 @@ import (
 	"unsafe"
 
 	"github.com/cockroachdb/pebble/internal/binfmt"
+	"github.com/cockroachdb/pebble/internal/invariants"
 	"github.com/cockroachdb/pebble/internal/treeprinter"
 )
 
@@ -117,6 +118,7 @@ func (b *RawBytes) slice(start, end uint32) []byte {
 
 // At returns the []byte at index i. The returned slice should not be mutated.
 func (b RawBytes) At(i int) []byte {
+	invariants.CheckBounds(i, b.slices)
 	return b.slice(b.offsets.At2(i))
 }