diff --git a/cmd/datdump/main.go b/cmd/datdump/main.go index d3629100..518c5d8a 100644 --- a/cmd/datdump/main.go +++ b/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 } diff --git a/main.go b/main.go index a646b47f..2f5fad39 100644 --- a/main.go +++ b/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 - finalMap map[string][]*Entry - cirIncMap map[string]bool + parsedListByName map[string]*ParsedList } 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) @@ -237,7 +240,8 @@ func parseEntry(typ, rule string) (*Entry, []string, error) { } } - slices.Sort(entry.Attrs) // Sort attributes + 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] == '-' { - battr := attr[1:] + if battr, ok := strings.CutPrefix(attr, "-"); ok { 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} - p.plMap[name] = pl + pl = new(ParsedList) + 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 { - pl, ok := p.plMap[plname] +// resolveList resolves the inclusions of the named list and returns it. +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] { - return fmt.Errorf("circular inclusion in: %q", plname) + if pl.Resolving { + return nil, fmt.Errorf("circular inclusion in: %q", plname) } - p.cirIncMap[plname] = true - defer delete(p.cirIncMap, plname) + pl.Resolving = true + defer func() { pl.Resolving = false }() - roughMap := make(map[string]*Entry) // Avoid basic duplicates - for _, dentry := range pl.Entries { // Add direct entries - roughMap[dentry.Plain] = dentry + roughEntries := make(map[string]*Entry) // Avoid basic duplicates + for _, dentry := range pl.Entries { // Add direct entries + roughEntries[dentry.Plain] = dentry } for _, inc := range pl.Inclusions { // Add included entries - if err := p.resolveList(inc.Source); err != nil { - return fmt.Errorf("failed to resolve inclusion %q: %w", inc.Source, err) + ipl, err := p.resolveList(inc.Source) + 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 - processor := &Processor{plMap: make(map[string]*ParsedList)} + // Parse all lists in the data directory + 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 - processor.finalMap = make(map[string][]*Entry, len(processor.plMap)) - processor.cirIncMap = make(map[string]bool) - for plname := range processor.plMap { - if err := processor.resolveList(plname); err != nil { + // Resolve the inclusions of all lists + for plname := range processor.parsedListByName { + 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)] - if !exist { - fmt.Printf("[Warn] list %q does not exist\n", epList) + pl, exist := processor.parsedListByName[strings.ToUpper(epList)] + if !exist || len(pl.FinalEntries) == 0 { + 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), - SiteIdx: make(map[string]int, sitesCount), + Sites: make([]*router.GeoSite, 0, listsCount), + SiteIdx: make(map[string]int, listsCount), } - for siteName, siteEntries := range processor.finalMap { - gs.Sites = append(gs.Sites, makeProtoList(siteName, siteEntries)) + for siteName, pl := range processor.parsedListByName { + 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 } diff --git a/main_test.go b/main_test.go new file mode 100644 index 00000000..c0b7f017 --- /dev/null +++ b/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) + } +}