Skip to content

Commit

Permalink
Fix: use valid country code as fallback when processing MaxMind data
Browse files Browse the repository at this point in the history
  • Loading branch information
Loyalsoldier committed Aug 3, 2024
1 parent efc3c0c commit ac37ef6
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 11 deletions.
17 changes: 14 additions & 3 deletions plugin/maxmind/country_csv.go
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ func (g *geoLite2CountryCSV) Input(container lib.Container) (lib.Container, erro
return nil, err
}

entries := make(map[string]*lib.Entry)
entries := make(map[string]*lib.Entry, 300)

if g.IPv4File != "" {
if err := g.process(g.IPv4File, ccMap, entries); err != nil {
Expand Down Expand Up @@ -174,7 +174,7 @@ func (g *geoLite2CountryCSV) process(file string, ccMap map[string]string, entri
return errors.New("country code list must be specified")
}
if entries == nil {
entries = make(map[string]*lib.Entry)
entries = make(map[string]*lib.Entry, 300)
}

fReader, err := os.Open(file)
Expand All @@ -198,7 +198,18 @@ func (g *geoLite2CountryCSV) process(file string, ccMap map[string]string, entri
}

for _, line := range lines[1:] {
ccID := strings.TrimSpace(line[1])
ccID := ""
switch {
case strings.TrimSpace(line[1]) != "":
ccID = strings.TrimSpace(line[1])
case strings.TrimSpace(line[2]) != "":
ccID = strings.TrimSpace(line[2])
case strings.TrimSpace(line[3]) != "":
ccID = strings.TrimSpace(line[3])
default:
continue
}

if countryCode, found := ccMap[ccID]; found {
if len(wantList) > 0 && !wantList[countryCode] {
continue
Expand Down
33 changes: 25 additions & 8 deletions plugin/maxmind/mmdb_in.go
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ func (g *maxmindMMDBIn) Input(container lib.Container) (lib.Container, error) {
return nil, err
}

entries := make(map[string]*lib.Entry)
entries := make(map[string]*lib.Entry, 300)
err = g.generateEntries(content, entries)
if err != nil {
return nil, err
Expand Down Expand Up @@ -146,20 +146,37 @@ func (g *maxmindMMDBIn) generateEntries(content []byte, entries map[string]*lib.
}
defer db.Close()

record := struct {
Country struct {
IsoCode string `maxminddb:"iso_code"`
} `maxminddb:"country"`
}{}

networks := db.Networks(maxminddb.SkipAliasedNetworks)
for networks.Next() {
record := struct {
Country struct {
IsoCode string `maxminddb:"iso_code"`
} `maxminddb:"country"`
RegisteredCountry struct {
IsoCode string `maxminddb:"iso_code"`
} `maxminddb:"registered_country"`
RepresentedCountry struct {
IsoCode string `maxminddb:"iso_code"`
} `maxminddb:"represented_country"`
}{}

subnet, err := networks.Network(&record)
if err != nil {
continue
}

name := strings.ToUpper(record.Country.IsoCode)
name := ""
switch {
case strings.TrimSpace(record.Country.IsoCode) != "":
name = strings.ToUpper(strings.TrimSpace(record.Country.IsoCode))
case strings.TrimSpace(record.RegisteredCountry.IsoCode) != "":
name = strings.ToUpper(strings.TrimSpace(record.RegisteredCountry.IsoCode))
case strings.TrimSpace(record.RepresentedCountry.IsoCode) != "":
name = strings.ToUpper(strings.TrimSpace(record.RepresentedCountry.IsoCode))
default:
continue
}

entry, found := entries[name]
if !found {
entry = lib.NewEntry(name)
Expand Down

0 comments on commit ac37ef6

Please sign in to comment.