From 58bccba2549bdb144e6c668a5aa882044e731710 Mon Sep 17 00:00:00 2001
From: Loyalsoldier <10487845+Loyalsoldier@users.noreply.github.com>
Date: Sun, 4 Aug 2024 05:47:07 +0800
Subject: [PATCH] Fix: use valid country code as fallback when processing
 MaxMind data

---
 plugin/maxmind/country_csv.go | 17 ++++++++++++++---
 plugin/maxmind/mmdb_in.go     | 33 +++++++++++++++++++++++++--------
 2 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/plugin/maxmind/country_csv.go b/plugin/maxmind/country_csv.go
index c5b276189..9c703c04e 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 7952af4cd..3c25931de 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)