diff --git a/plugin/maxmind/country_csv.go b/plugin/maxmind/country_csv.go index aea07a7313c..3511b0c7957 100644 --- a/plugin/maxmind/country_csv.go +++ b/plugin/maxmind/country_csv.go @@ -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 { @@ -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) @@ -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 diff --git a/plugin/maxmind/mmdb_in.go b/plugin/maxmind/mmdb_in.go index 3cf2f9a6298..065f5601858 100644 --- a/plugin/maxmind/mmdb_in.go +++ b/plugin/maxmind/mmdb_in.go @@ -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 @@ -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)