diff --git a/internal/xray/geodata/geodata.go b/internal/xray/geodata/geodata.go index c6ef6ae3b..caf802fe2 100644 --- a/internal/xray/geodata/geodata.go +++ b/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) { diff --git a/internal/xray/geodata/geodata_test.go b/internal/xray/geodata/geodata_test.go index 35ac99ac9..5582b722f 100644 --- a/internal/xray/geodata/geodata_test.go +++ b/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) diff --git a/internal/xray/geodata/reader.go b/internal/xray/geodata/reader.go index eaaa1bf4a..eed655f15 100644 --- a/internal/xray/geodata/reader.go +++ b/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:]