Refine: read large CSV file line by line

1 files changed, 38 insertions, 20 deletions

diff --git a/plugin/maxmind/country_csv.go b/plugin/maxmind/country_csv.go
index 3511b0c7..c729467a 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", typeASNCSV, 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)