index: surface corrupt rune and match offsets (#1108)

Malformed rune mappings previously underflowed into huge offsets, while
defensive bounds checks risked turning the same corruption into an
invisible non-match. Propagate the first per-document corruption error
through Search so callers can distinguish incomplete results from a
valid zero-match response.

Note: zoekt already handles panics in search as a corrupt index, and we
mark as part of the search statistics that we crashed a shard. This
commit makes it so we are far more defensive.

Author: keegancsmith

index/bits.go

@@ -57,9 +57,9 @@ func toLower(in []byte) []byte {
 }
 
 // compare 'lower' and 'mixed', where lower is the needle. 'mixed' may
-// be larger than 'lower'. Returns whether there was a match, and if
-// yes, the byte size of the match.
-func caseFoldingEqualsRunes(lower, mixed []byte) (int, bool) {
+// be larger than 'lower'. Returns the byte size of the match, whether there was
+// a match, and whether mixed ended before lower.
+func caseFoldingEqualsRunes(lower, mixed []byte) (int, bool, bool) {
 	matchTotal := 0
 	for len(lower) > 0 && len(mixed) > 0 {
 		lb := lower[0]
@@ -70,7 +70,7 @@ func caseFoldingEqualsRunes(lower, mixed []byte) (int, bool) {
 				mb |= 0x20
 			}
 			if lb != mb {
-				return 0, false
+				return 0, false, !hasEnoughRunes(mixed, lower)
 			}
 			lower = lower[1:]
 			mixed = mixed[1:]
@@ -86,11 +86,37 @@ func caseFoldingEqualsRunes(lower, mixed []byte) (int, bool) {
 		matchTotal += msz
 
 		if lr != unicode.ToLower(mr) {
-			return 0, false
+			return 0, false, !hasEnoughRunes(mixed, lower)
 		}
 	}
 
-	return matchTotal, len(lower) == 0
+	return matchTotal, len(lower) == 0, len(lower) > 0
+}
+
+func hasEnoughRunes(mixed, lower []byte) bool {
+	// A UTF-8 rune occupies at most UTFMax bytes, and lower cannot contain
+	// more runes than bytes. A sufficiently long mixed suffix therefore proves
+	// the candidate span fits without walking either slice. In practice this
+	// keeps ordinary mismatches constant-time and reserves the exact scan below
+	// for candidates close enough to the document end to be truncated.
+	if len(mixed)/utf8.UTFMax >= len(lower) {
+		return true
+	}
+
+	// Compare only rune availability, not values: the caller already knows the
+	// candidate mismatches. Reaching the end of mixed first means the candidate's
+	// expected rune span crosses the document boundary, which is a corrupt index
+	// invariant rather than an ordinary non-match.
+	for len(lower) > 0 {
+		if len(mixed) == 0 {
+			return false
+		}
+		_, sz := utf8.DecodeRune(lower)
+		lower = lower[sz:]
+		_, sz = utf8.DecodeRune(mixed)
+		mixed = mixed[sz:]
+	}
+	return true
 }
 
 type ngram uint64
@@ -394,12 +420,12 @@ func makeRuneOffsetMap(off []uint32) runeOffsetMap {
 // runes to traverse, given the granularity of runeOffsetFrequency.
 //
 // It does this by finding the nearest point to interpolate from in the map.
-func (m runeOffsetMap) lookup(runeOffset uint32) (uint32, uint32) {
+func (m runeOffsetMap) lookup(runeOffset uint32) (uint64, uint32) {
 	left := runeOffset % runeOffsetFrequency
 	runeOffset -= left
 	slen := len(m)
 	if slen == 0 {
-		return runeOffset, left
+		return uint64(runeOffset), left
 	}
 	// sort.Search finds the *first* index for which the predicate is true,
 	// but we want to find the *last* index for which the predicate is true.
@@ -410,9 +436,9 @@ func (m runeOffsetMap) lookup(runeOffset uint32) (uint32, uint32) {
 	idx = slen - 1 - idx
 	// idx is now in the range [-1, len(m))-- -1 indicates that the offset is smaller
 	// than the first entry in the map, so no correction is necessary.
-	byteOff := runeOffset
+	byteOff := uint64(runeOffset)
 	if idx >= 0 {
-		byteOff = m[idx].byteOffset + runeOffset - m[idx].runeOffset
+		byteOff = uint64(m[idx].byteOffset) + uint64(runeOffset) - uint64(m[idx].runeOffset)
 	}
 	return byteOff, left
 }

index/bits_test.go

@@ -224,7 +224,7 @@ func TestCondenseRuneOffsets(t *testing.T) {
 		for j, byteOffset := range tc.arr {
 			runeOffset := uint32(j * runeOffsetFrequency)
 			gotByteOffset, _ := got.lookup(runeOffset)
-			if gotByteOffset != byteOffset {
+			if gotByteOffset != uint64(byteOffset) {
 				t.Errorf("#%d: lookup(%v) got %v, want %v", i, runeOffset, gotByteOffset, byteOffset)
 			}
 		}
@@ -247,7 +247,7 @@ func TestRuneOffsetLookup(t *testing.T) {
 		if gotLeft != tc.wantLeft {
 			t.Errorf("#%d: got left=%v, want left=%v", i, gotLeft, tc.wantLeft)
 		}
-		if gotOff != tc.wantOff {
+		if gotOff != uint64(tc.wantOff) {
 			t.Errorf("#%d: got off=%v, want off=%v", i, gotOff, tc.wantOff)
 		}
 	}
@@ -257,7 +257,7 @@ func TestRuneOffsetLookup(t *testing.T) {
 	wanted := []uint32{0, 0, 0, 105, 105, 105, 210, 210, 310, 310, 430, 430, 530, 630}
 	for i, v := range inputs {
 		got, _ := m.lookup(v)
-		if got != wanted[i] {
+		if got != uint64(wanted[i]) {
 			t.Errorf("got off=%v, want off=%v for map=%v", got, wanted[i], m)
 		}
 	}

index/case_folding_bench_test.go

@@ -25,7 +25,7 @@ func TestCaseFoldingEqualsRunes(t *testing.T) {
 		{"äbč", "ÄBČ", true, 5}, // 'ä' (2 bytes), 'b' (1 byte), 'č' (2 bytes)
 		{"äbč", "ÄBX", false, 0},
 	} {
-		sz, ok := caseFoldingEqualsRunes([]byte(tc.lower), []byte(tc.mixed))
+		sz, ok, _ := caseFoldingEqualsRunes([]byte(tc.lower), []byte(tc.mixed))
 		if ok != tc.wantMatch || sz != tc.wantSz {
 			t.Errorf("caseFoldingEqualsRunes(%q, %q): got (%d, %t), want (%d, %t)",
 				tc.lower, tc.mixed, sz, ok, tc.wantSz, tc.wantMatch)
@@ -45,7 +45,7 @@ func BenchmarkCaseFoldingEqualsRunes(b *testing.B) {
 	b.Run("ASCII", func(b *testing.B) {
 		b.ReportAllocs()
 		for i := 0; i < b.N; i++ {
-			sz, ok := caseFoldingEqualsRunes(asciiLower, asciiMixed)
+			sz, ok, _ := caseFoldingEqualsRunes(asciiLower, asciiMixed)
 			if !ok || sz != len(asciiMixed) {
 				b.Fatalf("bad match: %d, %t", sz, ok)
 			}
@@ -55,7 +55,7 @@ func BenchmarkCaseFoldingEqualsRunes(b *testing.B) {
 	b.Run("Unicode", func(b *testing.B) {
 		b.ReportAllocs()
 		for i := 0; i < b.N; i++ {
-			sz, ok := caseFoldingEqualsRunes(unicodeLower, unicodeMixed)
+			sz, ok, _ := caseFoldingEqualsRunes(unicodeLower, unicodeMixed)
 			if !ok || sz != len(unicodeMixed) {
 				b.Fatalf("bad match: %d, %t", sz, ok)
 			}

index/contentprovider.go

@@ -16,6 +16,7 @@ package index
 
 import (
 	"bytes"
+	"fmt"
 	"log"
 	"path"
 	"slices"
@@ -36,7 +37,6 @@ type contentProvider struct {
 	stats *zoekt.Stats
 
 	// mutable
-	err      error
 	idx      uint32
 	_data    []byte
 	_nl      []uint32
@@ -58,10 +58,34 @@ func (p *contentProvider) setDocument(docID uint32) {
 	p._data = nil
 }
 
+// panicCorrupt stops a search as soon as it detects a corrupt shard invariant.
+// searchOneShard recovers the panic at the shard boundary, logs the error with
+// the query and stack trace, and increments Stats.Crashes. Callers therefore get
+// partial results with an explicit crashed-shard signal rather than a silent
+// non-match.
+func (p *contentProvider) panicCorrupt(err error) {
+	shard := "unknown"
+	if p.id.file != nil {
+		shard = p.id.file.Name()
+	}
+	repo := "unknown"
+	if p.idx < uint32(len(p.id.repos)) {
+		repoID := p.id.repos[p.idx]
+		if int(repoID) < len(p.id.repoMetaData) {
+			repo = p.id.repoMetaData[repoID].Name
+		}
+	}
+	panic(fmt.Errorf("corrupt shard %q while searching repository %q, document %d: %w", shard, repo, p.idx, err))
+}
+
 func (p *contentProvider) docSections() []DocumentSection {
 	if p._sects == nil {
 		var sz uint32
-		p._sects, sz, p.err = p.id.readDocSections(p.idx, p._sectBuf)
+		var err error
+		p._sects, sz, err = p.id.readDocSections(p.idx, p._sectBuf)
+		if err != nil {
+			p.panicCorrupt(fmt.Errorf("reading document sections: %w", err))
+		}
 		p.stats.ContentBytesLoaded += int64(sz)
 		p._sectBuf = p._sects
 	}
@@ -71,7 +95,11 @@ func (p *contentProvider) docSections() []DocumentSection {
 func (p *contentProvider) newlines() newlines {
 	if p._nl == nil {
 		var sz uint32
-		p._nl, sz, p.err = p.id.readNewlines(p.idx, p._nlBuf)
+		var err error
+		p._nl, sz, err = p.id.readNewlines(p.idx, p._nlBuf)
+		if err != nil {
+			p.panicCorrupt(fmt.Errorf("reading newline offsets: %w", err))
+		}
 		p._nlBuf = p._nl
 		p.stats.ContentBytesLoaded += int64(sz)
 	}
@@ -84,7 +112,11 @@ func (p *contentProvider) data(fileName bool) []byte {
 	}
 
 	if p._data == nil {
-		p._data, p.err = p.id.readContents(p.idx)
+		var err error
+		p._data, err = p.id.readContents(p.idx)
+		if err != nil {
+			p.panicCorrupt(fmt.Errorf("reading content: %w", err))
+		}
 		p.stats.FilesLoaded++
 		p.stats.ContentBytesLoaded += int64(len(p._data))
 	}
@@ -95,45 +127,86 @@ func (p *contentProvider) data(fileName bool) []byte {
 // runes (relative to document start). If filename is set, the corpus
 // is the set of filenames, with the document being the name itself.
 func (p *contentProvider) findOffset(filename bool, r uint32) uint32 {
-	if p.id.metaData.PlainASCII {
-		return r
-	}
-
-	sample := p.id.runeOffsets
-	runeEnds := p.id.fileEndRunes
-	fileStartByte := p.id.boundaries[p.idx]
+	var sample runeOffsetMap
+	var runeEnds []uint32
+	var fileStartByte, fileEndByte uint32
+	kind := "content"
 	if filename {
 		sample = p.id.fileNameRuneOffsets
 		runeEnds = p.id.fileNameEndRunes
 		fileStartByte = p.id.fileNameIndex[p.idx]
+		fileEndByte = p.id.fileNameIndex[p.idx+1]
+		kind = "filename"
+	} else {
+		sample = p.id.runeOffsets
+		runeEnds = p.id.fileEndRunes
+		fileStartByte = p.id.boundaries[p.idx]
+		fileEndByte = p.id.boundaries[p.idx+1]
 	}
 
-	absR := r
+	if p.id.metaData.PlainASCII {
+		if r > fileEndByte-fileStartByte {
+			p.panicCorrupt(fmt.Errorf("%s rune offset %d is after file size %d", kind, r, fileEndByte-fileStartByte))
+			return 0
+		}
+		return r
+	}
+
+	absR64 := uint64(r)
 	if p.idx > 0 {
-		absR += runeEnds[p.idx-1]
+		absR64 += uint64(runeEnds[p.idx-1])
+	}
+	if absR64 > uint64(^uint32(0)) {
+		p.panicCorrupt(fmt.Errorf("%s rune offset %d overflows the corpus rune offset", kind, r))
+		return 0
 	}
+	absR := uint32(absR64)
 
 	byteOff, left := sample.lookup(absR)
 
 	var data []byte
 
 	if filename {
-		data = p.id.fileNameContent[byteOff:]
+		if byteOff > uint64(len(p.id.fileNameContent)) {
+			p.panicCorrupt(fmt.Errorf("filename rune offset %d maps to byte offset %d past filename data size %d", absR, byteOff, len(p.id.fileNameContent)))
+			return 0
+		}
+		data = p.id.fileNameContent[uint32(byteOff):]
 	} else {
-		data, p.err = p.id.readContentSlice(byteOff, 3*runeOffsetFrequency)
-		if p.err != nil {
+		corpusEnd := p.id.boundaries[len(p.id.boundaries)-1]
+		if byteOff > uint64(corpusEnd) {
+			p.panicCorrupt(fmt.Errorf("content rune offset %d maps to byte offset %d past content data size %d", absR, byteOff, corpusEnd))
+			return 0
+		}
+		var err error
+		data, err = p.id.readContentSlice(uint32(byteOff), 3*runeOffsetFrequency)
+		if err != nil {
+			p.panicCorrupt(fmt.Errorf("content rune offset %d cannot load bytes at offset %d: %w", absR, byteOff, err))
 			return 0
 		}
 	}
 	for left > 0 {
+		if len(data) == 0 {
+			p.panicCorrupt(fmt.Errorf("%s rune offset %d has no decode bytes at byte offset %d", kind, absR, byteOff))
+			return 0
+		}
 		_, sz := utf8.DecodeRune(data)
-		byteOff += uint32(sz)
+		byteOff += uint64(sz)
 		data = data[sz:]
 		left--
 	}
 
-	byteOff -= fileStartByte
-	return byteOff
+	if byteOff < uint64(fileStartByte) {
+		p.panicCorrupt(fmt.Errorf("%s rune offset %d maps to byte offset %d before file start %d", kind, absR, byteOff, fileStartByte))
+		return 0
+	}
+	if byteOff > uint64(fileEndByte) {
+		p.panicCorrupt(fmt.Errorf("%s rune offset %d maps to byte offset %d after file end %d", kind, absR, byteOff, fileEndByte))
+		return 0
+	}
+
+	byteOff -= uint64(fileStartByte)
+	return uint32(byteOff)
 }
 
 // fillMatches converts the internal candidateMatch slice into our API's LineMatch.