Fix panic and rule loss in list parsing/resolution (#3905)

* Fix panic on empty inclusion attribute, rule loss in selective inclusion and other bugs

* Exercise domain/list-name validators in tests

* Consolidate parsed list state into ParsedList and simplify inclusion parsing

Co-authored-by: Loyalsoldier <10487845+Loyalsoldier@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: Loyalsoldier <10487845+Loyalsoldier@users.noreply.github.com>

cmd/datdump/main.go

@@ -143,20 +143,26 @@ func run() error {
 		exportListSlice = []string{"_all_"}
 	}
 
+	failedCount := 0
 	for _, eplistname := range exportListSlice {
 		if strings.EqualFold(eplistname, "_all_") {
 			if err := exportAll(filepath.Base(*inputData)+"_plain.yml", geoSites); err != nil {
 				fmt.Printf("[Error] failed to exportAll: %v\n", err)
+				failedCount++
 				continue
 			}
 		} else {
 			if err := exportSite(eplistname, geoSites); err != nil {
 				fmt.Printf("[Error] failed to exportSite: %v\n", err)
+				failedCount++
 				continue
 			}
 		}
 		fmt.Printf("list: %q has been exported successfully\n", eplistname)
 	}
+	if failedCount > 0 {
+		return fmt.Errorf("%d list(s) failed to be exported", failedCount)
+	}
 	return nil
 }
 

main.go

@@ -38,15 +38,17 @@ type Inclusion struct {
 }
 
 type ParsedList struct {
-	Resolved   bool
 	Inclusions []*Inclusion
-	Entries    []*Entry
+	Entries    []*Entry // Entries parsed from the list itself
+	// The fields below are filled in by resolveList
+	Resolving    bool
+	Resolved     bool
+	RoughEntries map[string]*Entry // Deduplicated direct and included entries
+	FinalEntries []*Entry          // Sorted entries without redundant subdomains
 }
 
 type Processor struct {
-	plMap     map[string]*ParsedList
+	parsedListByName map[string]*ParsedList
-	finalMap  map[string][]*Entry
-	cirIncMap map[string]bool
 }
 
 type GeoSites struct {
@@ -136,11 +138,12 @@ func (gs *GeoSites) assembleDat(task DatTask) error {
 				return fmt.Errorf("list %q not found for allowlist task", list)
 			}
 		}
+		slices.Sort(allowedIdxes)
+		allowedIdxes = slices.Compact(allowedIdxes) // Avoid duplicated lists
 		allowedlen := len(allowedIdxes)
 		if allowedlen == 0 {
 			return fmt.Errorf("allowlist needs at least one valid list")
 		}
-		slices.Sort(allowedIdxes)
 		geoSiteList.Entry = make([]*router.GeoSite, allowedlen)
 		for i, idx := range allowedIdxes {
 			geoSiteList.Entry[i] = gs.Sites[idx]
@@ -151,12 +154,12 @@ func (gs *GeoSites) assembleDat(task DatTask) error {
 			if idx, ok := gs.SiteIdx[strings.ToUpper(list)]; ok {
 				deniedMap[idx] = true
 			} else {
-				fmt.Printf("[Warn] list %q not found in denylist task %q", list, task.Name)
+				fmt.Printf("[Warn] list %q not found in denylist task %q\n", list, task.Name)
 			}
 		}
 		deniedlen := len(deniedMap)
 		if deniedlen == 0 {
-			fmt.Printf("[Warn] nothing to deny in task %q", task.Name)
+			fmt.Printf("[Warn] nothing to deny in task %q\n", task.Name)
 			geoSiteList.Entry = gs.Sites
 		} else {
 			geoSiteList.Entry = make([]*router.GeoSite, 0, len(gs.Sites)-deniedlen)
@@ -238,6 +241,7 @@ func parseEntry(typ, rule string) (*Entry, []string, error) {
 	}
 
 	slices.Sort(entry.Attrs)                  // Sort attributes
+	entry.Attrs = slices.Compact(entry.Attrs) // Remove duplicated attributes
 	// Formated plain entry: type:domain.tld:@attr1,@attr2
 	var plain strings.Builder
 	plain.Grow(plen)
@@ -272,8 +276,7 @@ func parseInclusion(rule string) (*Inclusion, error) {
 		switch part[0] {
 		case '@':
 			attr := strings.ToLower(part[1:])
-			if attr[0] == '-' {
+			if battr, ok := strings.CutPrefix(attr, "-"); ok {
-				battr := attr[1:]
 				if !validateAttrChars(battr) {
 					return inc, fmt.Errorf("invalid ban attribute: %q", battr)
 				}
@@ -336,10 +339,10 @@ func validateSiteName(name string) bool {
 }
 
 func (p *Processor) getOrCreateParsedList(name string) *ParsedList {
-	pl, exist := p.plMap[name]
+	pl, exist := p.parsedListByName[name]
 	if !exist {
-		pl = &ParsedList{Resolved: false}
+		pl = new(ParsedList)
-		p.plMap[name] = pl
+		p.parsedListByName[name] = pl
 	}
 	return pl
 }
@@ -457,49 +460,55 @@ func polishList(roughMap map[string]*Entry) []*Entry {
 	return finalList
 }
 
-func (p *Processor) resolveList(plname string) error {
+// resolveList resolves the inclusions of the named list and returns it.
-	pl, ok := p.plMap[plname]
+func (p *Processor) resolveList(plname string) (*ParsedList, error) {
+	pl, ok := p.parsedListByName[plname]
 	if !ok {
-		return fmt.Errorf("list %q not found", plname)
+		return nil, fmt.Errorf("list %q not found", plname)
 	}
 	if pl.Resolved {
-		return nil
+		return pl, nil
 	}
-	if p.cirIncMap[plname] {
+	if pl.Resolving {
-		return fmt.Errorf("circular inclusion in: %q", plname)
+		return nil, fmt.Errorf("circular inclusion in: %q", plname)
 	}
-	p.cirIncMap[plname] = true
+	pl.Resolving = true
-	defer delete(p.cirIncMap, plname)
+	defer func() { pl.Resolving = false }()
 
-	roughMap := make(map[string]*Entry) // Avoid basic duplicates
+	roughEntries := make(map[string]*Entry) // Avoid basic duplicates
 	for _, dentry := range pl.Entries {     // Add direct entries
-		roughMap[dentry.Plain] = dentry
+		roughEntries[dentry.Plain] = dentry
 	}
 	for _, inc := range pl.Inclusions { // Add included entries
-		if err := p.resolveList(inc.Source); err != nil {
+		ipl, err := p.resolveList(inc.Source)
-			return fmt.Errorf("failed to resolve inclusion %q: %w", inc.Source, err)
+		if err != nil {
+			return nil, fmt.Errorf("failed to resolve inclusion %q: %w", inc.Source, err)
 		}
 		isFullInc := len(inc.MustAttrs) == 0 && len(inc.BanAttrs) == 0
-		for _, ientry := range p.finalMap[inc.Source] {
+		// Filter the unpolished entries of the source list, otherwise selective
+		// inclusion would lose rules that have been pruned in the source list as
+		// redundant subdomains of a parent rule which is filtered out here.
+		for _, ientry := range ipl.RoughEntries {
 			if isFullInc || isMatchAttrFilters(ientry, inc) {
-				roughMap[ientry.Plain] = ientry
+				roughEntries[ientry.Plain] = ientry
 			}
 		}
 	}
-	if len(roughMap) == 0 {
+	pl.RoughEntries = roughEntries
+	if len(roughEntries) == 0 {
 		fmt.Printf("[Warn] ignore empty list %q\n", plname)
 	} else {
-		p.finalMap[plname] = polishList(roughMap)
+		pl.FinalEntries = polishList(roughEntries)
 	}
 	pl.Resolved = true
-	return nil
+	return pl, nil
 }
 
 func run() error {
 	fmt.Printf("using domain lists data in %q\n", *dataPath)
 
-	// Generate plMap
+	// Parse all lists in the data directory
-	processor := &Processor{plMap: make(map[string]*ParsedList)}
+	processor := &Processor{parsedListByName: make(map[string]*ParsedList)}
 	err := filepath.WalkDir(*dataPath, func(path string, d os.DirEntry, err error) error {
 		if err != nil {
 			return err
@@ -516,30 +525,29 @@ func run() error {
 	if err != nil {
 		return fmt.Errorf("failed to loadData: %w", err)
 	}
-	// Generate finalMap
+	// Resolve the inclusions of all lists
-	processor.finalMap = make(map[string][]*Entry, len(processor.plMap))
+	for plname := range processor.parsedListByName {
-	processor.cirIncMap = make(map[string]bool)
+		if _, err := processor.resolveList(plname); err != nil {
-	for plname := range processor.plMap {
-		if err := processor.resolveList(plname); err != nil {
 			return fmt.Errorf("failed to resolveList %q: %w", plname, err)
 		}
 	}
-	processor.plMap = nil
 
 	// Make sure output directory exists
 	if err := os.MkdirAll(*outputDir, 0755); err != nil {
 		return fmt.Errorf("failed to create output directory: %w", err)
 	}
 	// Export plaintext lists
+	failedCount := 0
 	for rawEpList := range strings.SplitSeq(*exportLists, ",") {
 		if epList := strings.TrimSpace(rawEpList); epList != "" {
-			entries, exist := processor.finalMap[strings.ToUpper(epList)]
+			pl, exist := processor.parsedListByName[strings.ToUpper(epList)]
-			if !exist {
+			if !exist || len(pl.FinalEntries) == 0 {
-				fmt.Printf("[Warn] list %q does not exist\n", epList)
+				fmt.Printf("[Warn] list %q does not exist or is empty\n", epList)
 				continue
 			}
-			if err := writePlainList(epList, entries); err != nil {
+			if err := writePlainList(epList, pl.FinalEntries); err != nil {
 				fmt.Printf("[Error] failed to write list %q: %v\n", epList, err)
+				failedCount++
 				continue
 			}
 			fmt.Printf("list %q has been generated successfully\n", epList)
@@ -547,20 +555,22 @@ func run() error {
 	}
 
 	// Generate proto sites
-	sitesCount := len(processor.finalMap)
+	listsCount := len(processor.parsedListByName)
 	gs := &GeoSites{
-		Sites:   make([]*router.GeoSite, 0, sitesCount),
+		Sites:   make([]*router.GeoSite, 0, listsCount),
-		SiteIdx: make(map[string]int, sitesCount),
+		SiteIdx: make(map[string]int, listsCount),
 	}
-	for siteName, siteEntries := range processor.finalMap {
+	for siteName, pl := range processor.parsedListByName {
-		gs.Sites = append(gs.Sites, makeProtoList(siteName, siteEntries))
+		if len(pl.FinalEntries) == 0 { // Skip empty lists
+			continue
+		}
+		gs.Sites = append(gs.Sites, makeProtoList(siteName, pl.FinalEntries))
 	}
-	processor = nil
 	// Sort proto sites so the generated file is reproducible
 	slices.SortFunc(gs.Sites, func(a, b *router.GeoSite) int {
 		return strings.Compare(a.CountryCode, b.CountryCode)
 	})
-	for i := range sitesCount {
+	for i := range gs.Sites {
 		gs.SiteIdx[gs.Sites[i].CountryCode] = i
 	}
 
@@ -577,9 +587,13 @@ func run() error {
 	}
 	for _, task := range tasks {
 		if err := gs.assembleDat(task); err != nil {
-			fmt.Printf("[Error] failed to assembleDat %q: %v", task.Name, err)
+			fmt.Printf("[Error] failed to assembleDat %q: %v\n", task.Name, err)
+			failedCount++
 		}
 	}
+	if failedCount > 0 {
+		return fmt.Errorf("%d output file(s) failed to be generated", failedCount)
+	}
 	return nil
 }
 

main_test.go

@@ -0,0 +1,205 @@
+package main
+
+import (
+	"os"
+	"path/filepath"
+	"slices"
+	"strings"
+	"testing"
+)
+
+func TestParseEntry(t *testing.T) {
+	testCases := []struct {
+		name      string
+		typ       string
+		rule      string
+		wantPlain string
+		wantAffs  []string
+		wantErr   bool
+	}{
+		{name: "domain", typ: "domain", rule: "Example.COM", wantPlain: "domain:example.com"},
+		{name: "sorted attrs", typ: "full", rule: "a.example.com @cn @ads", wantPlain: "full:a.example.com:@ads,@cn"},
+		{name: "duplicated attrs", typ: "domain", rule: "example.com @ads @ads", wantPlain: "domain:example.com:@ads"},
+		{name: "affiliations", typ: "domain", rule: "example.com &other @ads", wantPlain: "domain:example.com:@ads", wantAffs: []string{"OTHER"}},
+		{name: "regexp", typ: "regexp", rule: `^example\.com$`, wantPlain: `regexp:^example\.com$`},
+		{name: "invalid regexp", typ: "regexp", rule: "^example(", wantErr: true},
+		{name: "empty rule", typ: "domain", rule: "   ", wantErr: true},
+		{name: "unknown type", typ: "prefix", rule: "example.com", wantErr: true},
+		{name: "invalid domain", typ: "domain", rule: "exa_mple.com @ads", wantErr: true},
+		{name: "empty attr", typ: "domain", rule: "example.com @", wantErr: true},
+		{name: "empty affiliation", typ: "domain", rule: "example.com &", wantErr: true},
+		{name: "unknown field", typ: "domain", rule: "example.com ads", wantErr: true},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			entry, affs, err := parseEntry(tc.typ, tc.rule)
+			if tc.wantErr {
+				if err == nil {
+					t.Fatalf("parseEntry(%q, %q) = %q, want error", tc.typ, tc.rule, entry.Plain)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("parseEntry(%q, %q) got unexpected error: %v", tc.typ, tc.rule, err)
+			}
+			if entry.Plain != tc.wantPlain {
+				t.Errorf("parseEntry(%q, %q) = %q, want %q", tc.typ, tc.rule, entry.Plain, tc.wantPlain)
+			}
+			if len(affs) != len(tc.wantAffs) {
+				t.Fatalf("parseEntry(%q, %q) affiliations = %v, want %v", tc.typ, tc.rule, affs, tc.wantAffs)
+			}
+			for i, aff := range affs {
+				if aff != tc.wantAffs[i] {
+					t.Errorf("parseEntry(%q, %q) affiliations = %v, want %v", tc.typ, tc.rule, affs, tc.wantAffs)
+				}
+			}
+		})
+	}
+}
+
+func TestParseInclusion(t *testing.T) {
+	testCases := []struct {
+		name     string
+		rule     string
+		wantSrc  string
+		wantMust []string
+		wantBan  []string
+		wantErr  bool
+	}{
+		{name: "plain", rule: "other-list", wantSrc: "OTHER-LIST"},
+		{name: "filters", rule: "other @ads @-cn", wantSrc: "OTHER", wantMust: []string{"ads"}, wantBan: []string{"cn"}},
+		{name: "empty attr", rule: "other @", wantErr: true},
+		{name: "empty ban attr", rule: "other @-", wantErr: true},
+		{name: "empty rule", rule: " ", wantErr: true},
+		{name: "invalid name", rule: "other@list", wantErr: true},
+		{name: "affiliation", rule: "other &another", wantErr: true},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			inc, err := parseInclusion(tc.rule)
+			if tc.wantErr {
+				if err == nil {
+					t.Fatalf("parseInclusion(%q) = %+v, want error", tc.rule, inc)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("parseInclusion(%q) got unexpected error: %v", tc.rule, err)
+			}
+			if inc.Source != tc.wantSrc {
+				t.Errorf("parseInclusion(%q) source = %q, want %q", tc.rule, inc.Source, tc.wantSrc)
+			}
+			if len(inc.MustAttrs) != len(tc.wantMust) || len(inc.BanAttrs) != len(tc.wantBan) {
+				t.Fatalf("parseInclusion(%q) filters = %v/%v, want %v/%v", tc.rule, inc.MustAttrs, inc.BanAttrs, tc.wantMust, tc.wantBan)
+			}
+			for i, attr := range inc.MustAttrs {
+				if attr != tc.wantMust[i] {
+					t.Errorf("parseInclusion(%q) must attrs = %v, want %v", tc.rule, inc.MustAttrs, tc.wantMust)
+				}
+			}
+			for i, attr := range inc.BanAttrs {
+				if attr != tc.wantBan[i] {
+					t.Errorf("parseInclusion(%q) ban attrs = %v, want %v", tc.rule, inc.BanAttrs, tc.wantBan)
+				}
+			}
+		})
+	}
+}
+
+func TestPolishList(t *testing.T) {
+	rules := []struct{ typ, rule string }{
+		{"domain", "example.com"},
+		{"domain", "sub.example.com"},      // Redundant
+		{"full", "www.example.com"},        // Redundant
+		{"full", "example.com"},            // Redundant
+		{"full", "example.org"},            // Kept, no parent domain rule
+		{"domain", "ads.example.com @ads"}, // Kept, has attribute
+		{"keyword", "example"},
+	}
+	roughMap := make(map[string]*Entry, len(rules))
+	for _, r := range rules {
+		entry, _, err := parseEntry(r.typ, r.rule)
+		if err != nil {
+			t.Fatalf("parseEntry(%q, %q) got unexpected error: %v", r.typ, r.rule, err)
+		}
+		roughMap[entry.Plain] = entry
+	}
+	want := []string{"domain:ads.example.com:@ads", "domain:example.com", "full:example.org", "keyword:example"}
+	assertPlains(t, "polishList", polishList(roughMap), want)
+}
+
+// TestResolveSelectiveInclusion makes sure that selective inclusion does not
+// lose rules which are redundant in the source list only.
+func TestResolveSelectiveInclusion(t *testing.T) {
+	dataPath := t.TempDir()
+	files := map[string]string{
+		"source": "domain:example.com @cn\nfull:mail.example.com\ndomain:sub.example.com\ndomain:example.org @ads\n",
+		"banned": "include:source @-cn\n",
+		"must":   "include:source @ads\n",
+		"full":   "include:source\n",
+	}
+	for name, content := range files {
+		if err := os.WriteFile(filepath.Join(dataPath, name), []byte(content), 0644); err != nil {
+			t.Fatalf("failed to write test data %q: %v", name, err)
+		}
+	}
+
+	processor := &Processor{parsedListByName: make(map[string]*ParsedList)}
+	for name := range files {
+		if err := processor.loadData(strings.ToUpper(name), filepath.Join(dataPath, name)); err != nil {
+			t.Fatalf("loadData(%q) got unexpected error: %v", name, err)
+		}
+	}
+	for name := range files {
+		if _, err := processor.resolveList(strings.ToUpper(name)); err != nil {
+			t.Fatalf("resolveList(%q) got unexpected error: %v", name, err)
+		}
+	}
+
+	assertList(t, processor, "SOURCE", []string{"domain:example.com:@cn", "domain:example.org:@ads"})
+	assertList(t, processor, "FULL", []string{"domain:example.com:@cn", "domain:example.org:@ads"})
+	assertList(t, processor, "MUST", []string{"domain:example.org:@ads"})
+	assertList(t, processor, "BANNED", []string{"domain:example.org:@ads", "domain:sub.example.com", "full:mail.example.com"})
+}
+
+func TestResolveCircularInclusion(t *testing.T) {
+	dataPath := t.TempDir()
+	files := map[string]string{
+		"first":  "domain:example.com\ninclude:second\n",
+		"second": "include:first\n",
+	}
+	for name, content := range files {
+		if err := os.WriteFile(filepath.Join(dataPath, name), []byte(content), 0644); err != nil {
+			t.Fatalf("failed to write test data %q: %v", name, err)
+		}
+	}
+	processor := &Processor{parsedListByName: make(map[string]*ParsedList)}
+	for name := range files {
+		if err := processor.loadData(strings.ToUpper(name), filepath.Join(dataPath, name)); err != nil {
+			t.Fatalf("loadData(%q) got unexpected error: %v", name, err)
+		}
+	}
+	if _, err := processor.resolveList("FIRST"); err == nil {
+		t.Fatal("resolveList(\"FIRST\") = nil, want circular inclusion error")
+	}
+}
+
+func assertList(t *testing.T, p *Processor, name string, want []string) {
+	t.Helper()
+	pl, exist := p.parsedListByName[name]
+	if !exist {
+		t.Fatalf("list %q does not exist", name)
+	}
+	assertPlains(t, name, pl.FinalEntries, want)
+}
+
+func assertPlains(t *testing.T, name string, entries []*Entry, want []string) {
+	t.Helper()
+	got := make([]string, 0, len(entries))
+	for _, entry := range entries {
+		got = append(got, entry.Plain)
+	}
+	if !slices.Equal(got, want) {
+		t.Errorf("%s = %v, want %v", name, got, want)
+	}
+}