fix(geodata): refuse a database whose category index is too large to browse

The geofile browser indexes every category record and attribute key a
database holds, which cost the panel far more memory than the file's
bytes: a crafted file of a few megabytes full of empty categories or
attribute keys ran the panel out of memory. The file size cap bounded
nothing here.

The scan now counts records and attribute keys as it keeps them and
refuses a database past 65536 with "too large to browse", which the
cache remembers; the other layout is not retried for the same file.

internal/xray/geodata/geodata.go

@@ -30,9 +30,15 @@ const MaxFileSize int64 = 256 << 20
 // the caller asks for.
 const MaxPageSize = 500
 
+// MaxIndexSize caps the category records plus attribute keys one database may
+// keep in the index: each costs hundreds of bytes in memory for ten on disk.
+const MaxIndexSize = 1 << 16
+
 var (
 	// ErrFileTooLarge reports a database above MaxFileSize.
 	ErrFileTooLarge = errors.New("geodata file is too large to browse")
+	// ErrIndexTooLarge reports a database above MaxIndexSize.
+	ErrIndexTooLarge = errors.New("geodata category index is too large to browse")
 	// ErrInvalidName reports a file name that does not resolve to a .dat file
 	// directly inside the asset directory.
 	ErrInvalidName = errors.New("invalid geodata file name")
@@ -226,7 +232,8 @@ func (s *Store) index(name string) (*index, error) {
 // isPermanent reports whether an error will repeat for the same bytes, and is
 // therefore worth caching instead of re-deriving on every request.
 func isPermanent(err error) bool {
-	return errors.Is(err, ErrUnrecognized) || errors.Is(err, ErrInvalidName) || errors.Is(err, ErrFileTooLarge)
+	return errors.Is(err, ErrUnrecognized) || errors.Is(err, ErrInvalidName) || errors.Is(err, ErrFileTooLarge) ||
+		errors.Is(err, ErrIndexTooLarge)
 }
 
 func (s *Store) cachedIndex(key fileKey) (*index, bool) {

internal/xray/geodata/geodata_test.go

@@ -439,6 +439,66 @@ func TestCategorySplitOverManyRecordsIsIndexedInOnePass(t *testing.T) {
 	}
 }
 
+func TestFileAboveIndexLimitIsRejected(t *testing.T) {
+	rule := domain(xraygeodata.Domain_Full, "a.b")
+	categories := make([]*xraygeodata.GeoSite, MaxIndexSize+1)
+	records := make([]*xraygeodata.GeoSite, MaxIndexSize+1)
+	keys := make([]string, MaxIndexSize)
+	tagged := make([]*xraygeodata.Domain, MaxIndexSize)
+	for i := range categories {
+		categories[i] = site(strconv.Itoa(i))
+		records[i] = site("z", rule)
+	}
+	categories[0].Domain = append(categories[0].Domain, rule)
+	for i := range keys {
+		keys[i] = strconv.Itoa(i)
+		tagged[i] = domain(xraygeodata.Domain_Full, "a.b", keys[i])
+	}
+
+	tests := []struct {
+		name  string
+		sites []*xraygeodata.GeoSite
+		want  error
+	}{
+		{name: "categories at the limit", sites: categories[:MaxIndexSize]},
+		{name: "categories", sites: categories, want: ErrIndexTooLarge},
+		{name: "records of one category", sites: records, want: ErrIndexTooLarge},
+		{name: "attributes at the limit", sites: []*xraygeodata.GeoSite{site("z", tagged[1:]...)}},
+		{name: "attributes across rules", sites: []*xraygeodata.GeoSite{site("z", tagged...)}, want: ErrIndexTooLarge},
+		{
+			name:  "attributes of one rule",
+			sites: []*xraygeodata.GeoSite{site("z", domain(xraygeodata.Domain_Full, "a.b", keys...))},
+			want:  ErrIndexTooLarge,
+		},
+		{
+			name:  "one attribute repeated on one rule",
+			sites: []*xraygeodata.GeoSite{site("z", domain(xraygeodata.Domain_Full, "a.b", strings.Fields(strings.Repeat("k ", MaxIndexSize+1))...))},
+			want:  ErrIndexTooLarge,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			dir := t.TempDir()
+			writeSiteDB(t, dir, "geosite.dat", tt.sites...)
+			store := NewStore(dir)
+
+			if _, err := store.Categories("geosite.dat", "", 0, 10); !errors.Is(err, tt.want) {
+				t.Fatalf("error = %v, want %v", err, tt.want)
+			}
+			if len(store.indexes) != 1 {
+				t.Errorf("cached indexes = %d, want the outcome cached once", len(store.indexes))
+			}
+			files, err := store.ListFiles()
+			if err != nil {
+				t.Fatalf("ListFiles() error = %v", err)
+			}
+			if tt.want != nil && (len(files) != 1 || files[0].Error != tt.want.Error()) {
+				t.Errorf("ListFiles() = %+v, want the file listed with %q", files, tt.want)
+			}
+		})
+	}
+}
+
 func TestIndexCacheInvalidatedWhenFileChanges(t *testing.T) {
 	dir := t.TempDir()
 	sampleSiteDB(t, dir)

internal/xray/geodata/reader.go

@@ -57,6 +57,9 @@ func scanIndex(data []byte, kind GeoKind) (*categoryScan, error) {
 	// Parallel to scan.categories: a category split over many records adds each
 	// record's own attributes to its set and is sorted once, at the end.
 	var attributeSets []map[string]struct{}
+	// A record or attribute key costs the index far more than its bytes on disk, so
+	// MaxFileSize bounds nothing here: size counts what the scan has kept.
+	size := 0
 
 	err := eachListEntry(data, func(entry []byte, span byteSpan) error {
 		count := 0
@@ -73,6 +76,9 @@ func scanIndex(data []byte, kind GeoKind) (*categoryScan, error) {
 				for _, attr := range attrs {
 					attributes[attr] = struct{}{}
 				}
+				if len(attributes) > MaxIndexSize {
+					return ErrIndexTooLarge
+				}
 			} else {
 				_, ok, err := cidrBytes(payload)
 				if err != nil {
@@ -88,6 +94,9 @@ func scanIndex(data []byte, kind GeoKind) (*categoryScan, error) {
 		if err != nil || code == "" {
 			return err
 		}
+		if size += 1 + len(attributes); size > MaxIndexSize {
+			return ErrIndexTooLarge
+		}
 		scan.usable += count
 		scan.spans[code] = append(scan.spans[code], span)
 		if position, seen := byCode[code]; seen {
@@ -172,6 +181,10 @@ func detectKind(data []byte, name string) (GeoKind, *categoryScan, error) {
 	var firstErr error
 	for _, kind := range [...]GeoKind{first, second} {
 		scan, err := scanIndex(data, kind)
+		if errors.Is(err, ErrIndexTooLarge) {
+			// Retrying as the other layout would only spend the same memory again.
+			return "", nil, err
+		}
 		if err != nil {
 			if firstErr == nil {
 				firstErr = err
@@ -350,6 +363,9 @@ func domainValue(payload []byte) ([]byte, []string, error) {
 				return nil, nil, protowire.ParseError(size)
 			}
 			if key := attributeKey(raw); key != "" {
+				if len(attributes) == MaxIndexSize {
+					return nil, nil, ErrIndexTooLarge
+				}
 				attributes = append(attributes, key)
 			}
 			payload = payload[size:]