diff --git a/plugin/maxmind/country_csv.go b/plugin/maxmind/country_csv.go index 9c703c04..dfb36dd2 100644 --- a/plugin/maxmind/country_csv.go +++ b/plugin/maxmind/country_csv.go @@ -3,8 +3,8 @@ package maxmind import ( "encoding/csv" "encoding/json" - "errors" "fmt" + "io" "os" "path/filepath" "strings" @@ -159,6 +159,10 @@ func (g *geoLite2CountryCSV) getCountryCode() (map[string]string, error) { ccMap := make(map[string]string) for _, line := range lines[1:] { + if len(line) < 5 { + return nil, fmt.Errorf("❌ [type %s | action %s] invalid record: %v", typeCountryCSV, g.Action, line) + } + id := strings.TrimSpace(line[0]) countryCode := strings.TrimSpace(line[4]) if id == "" || countryCode == "" { @@ -166,17 +170,30 @@ func (g *geoLite2CountryCSV) getCountryCode() (map[string]string, error) { } ccMap[id] = strings.ToUpper(countryCode) } + + if len(ccMap) == 0 { + return nil, fmt.Errorf("❌ [type %s | action %s] invalid country code data", typeCountryCSV, g.Action) + } + return ccMap, nil } func (g *geoLite2CountryCSV) process(file string, ccMap map[string]string, entries map[string]*lib.Entry) error { if len(ccMap) == 0 { - return errors.New("country code list must be specified") + return fmt.Errorf("❌ [type %s | action %s] invalid country code data", typeCountryCSV, g.Action) } if entries == nil { entries = make(map[string]*lib.Entry, 300) } + // Filter want list + wantList := make(map[string]bool) + for _, want := range g.Want { + if want = strings.ToUpper(strings.TrimSpace(want)); want != "" { + wantList[want] = true + } + } + fReader, err := os.Open(file) if err != nil { return err @@ -184,28 +201,29 @@ func (g *geoLite2CountryCSV) process(file string, ccMap map[string]string, entri defer fReader.Close() reader := csv.NewReader(fReader) - lines, err := reader.ReadAll() - if err != nil { - return err - } + reader.Read() // skip header - // Filter want list - wantList := make(map[string]bool) - for _, want := range g.Want { - if want = strings.ToUpper(strings.TrimSpace(want)); want != "" { - wantList[want] = true + for { + record, err := reader.Read() + if err == io.EOF { + break + } + if err != nil { + return err + } + + if len(record) < 4 { + return fmt.Errorf("❌ [type %s | action %s] invalid record: %v", typeCountryCSV, g.Action, record) } - } - for _, line := range lines[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]) + case strings.TrimSpace(record[1]) != "": + ccID = strings.TrimSpace(record[1]) + case strings.TrimSpace(record[2]) != "": + ccID = strings.TrimSpace(record[2]) + case strings.TrimSpace(record[3]) != "": + ccID = strings.TrimSpace(record[3]) default: continue } @@ -214,7 +232,7 @@ func (g *geoLite2CountryCSV) process(file string, ccMap map[string]string, entri if len(wantList) > 0 && !wantList[countryCode] { continue } - cidrStr := strings.ToLower(strings.TrimSpace(line[0])) + cidrStr := strings.ToLower(strings.TrimSpace(record[0])) entry, found := entries[countryCode] if !found { entry = lib.NewEntry(countryCode)