Refine: wantedList in various formats

3 files changed, 54 insertions, 51 deletions

diff --git a/plugin/maxmind/country_csv.go b/plugin/maxmind/country_csv.go
index 15c80719..4aa02975 100644
--- a/plugin/maxmind/country_csv.go
+++ b/plugin/maxmind/country_csv.go
@@ -59,6 +59,14 @@ func newGeoLite2CountryCSV(action lib.Action, data json.RawMessage) (lib.InputCo
 		tmp.IPv6File = defaultCountryIPv6File
 	}
 
+	// Filter want list
+	wantList := make(map[string]bool)
+	for _, want := range tmp.Want {
+		if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
+			wantList[want] = true
+		}
+	}
+
 	return &geoLite2CountryCSV{
 		Type:            typeCountryCSV,
 		Action:          action,
@@ -66,7 +74,7 @@ func newGeoLite2CountryCSV(action lib.Action, data json.RawMessage) (lib.InputCo
 		CountryCodeFile: tmp.CountryCodeFile,
 		IPv4File:        tmp.IPv4File,
 		IPv6File:        tmp.IPv6File,
-		Want:            tmp.Want,
+		Want:            wantList,
 		OnlyIPType:      tmp.OnlyIPType,
 	}, nil
 }
@@ -78,7 +86,7 @@ type geoLite2CountryCSV struct {
 	CountryCodeFile string
 	IPv4File        string
 	IPv6File        string
-	Want            []string
+	Want            map[string]bool
 	OnlyIPType      lib.IPType
 }
 
@@ -100,7 +108,7 @@ func (g *geoLite2CountryCSV) Input(container lib.Container) (lib.Container, erro
 		return nil, err
 	}
 
-	entries := make(map[string]*lib.Entry, 300)
+	entries := make(map[string]*lib.Entry, len(ccMap))
 
 	if g.IPv4File != "" {
 		if err := g.process(g.IPv4File, ccMap, entries); err != nil {
@@ -164,11 +172,16 @@ func (g *geoLite2CountryCSV) getCountryCode() (map[string]string, error) {
 		}
 
 		id := strings.TrimSpace(line[0])
-		countryCode := strings.TrimSpace(line[4])
+		countryCode := strings.ToUpper(strings.TrimSpace(line[4]))
 		if id == "" || countryCode == "" {
 			continue
 		}
-		ccMap[id] = strings.ToUpper(countryCode)
+
+		if len(g.Want) > 0 && !g.Want[countryCode] {
+			continue
+		}
+
+		ccMap[id] = countryCode
 	}
 
 	if len(ccMap) == 0 {
@@ -183,15 +196,7 @@ func (g *geoLite2CountryCSV) process(file string, ccMap map[string]string, entri
 		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
-		}
+		entries = make(map[string]*lib.Entry, len(ccMap))
 	}
 
 	fReader, err := os.Open(file)
@@ -229,17 +234,16 @@ func (g *geoLite2CountryCSV) process(file string, ccMap map[string]string, entri
 		}
 
 		if countryCode, found := ccMap[ccID]; found {
-			if len(wantList) > 0 && !wantList[countryCode] {
-				continue
-			}
 			cidrStr := strings.ToLower(strings.TrimSpace(record[0]))
-			entry, found := entries[countryCode]
-			if !found {
+			entry, got := entries[countryCode]
+			if !got {
 				entry = lib.NewEntry(countryCode)
 			}
+
 			if err := entry.AddPrefix(cidrStr); err != nil {
 				return err
 			}
+
 			entries[countryCode] = entry
 		}
 	}
diff --git a/plugin/maxmind/mmdb_in.go b/plugin/maxmind/mmdb_in.go
index 065f5601..3ee1eacb 100644
--- a/plugin/maxmind/mmdb_in.go
+++ b/plugin/maxmind/mmdb_in.go
@@ -46,12 +46,20 @@ func newMaxmindMMDBIn(action lib.Action, data json.RawMessage) (lib.InputConvert
 		tmp.URI = defaultMMDBFile
 	}
 
+	// Filter want list
+	wantList := make(map[string]bool)
+	for _, want := range tmp.Want {
+		if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
+			wantList[want] = true
+		}
+	}
+
 	return &maxmindMMDBIn{
 		Type:        typeMaxmindMMDBIn,
 		Action:      action,
 		Description: descMaxmindMMDBIn,
 		URI:         tmp.URI,
-		Want:        tmp.Want,
+		Want:        wantList,
 		OnlyIPType:  tmp.OnlyIPType,
 	}, nil
 }
@@ -61,68 +69,55 @@ type maxmindMMDBIn struct {
 	Action      lib.Action
 	Description string
 	URI         string
-	Want        []string
+	Want        map[string]bool
 	OnlyIPType  lib.IPType
 }
 
-func (g *maxmindMMDBIn) GetType() string {
-	return g.Type
+func (m *maxmindMMDBIn) GetType() string {
+	return m.Type
 }
 
-func (g *maxmindMMDBIn) GetAction() lib.Action {
-	return g.Action
+func (m *maxmindMMDBIn) GetAction() lib.Action {
+	return m.Action
 }
 
-func (g *maxmindMMDBIn) GetDescription() string {
-	return g.Description
+func (m *maxmindMMDBIn) GetDescription() string {
+	return m.Description
 }
 
-func (g *maxmindMMDBIn) Input(container lib.Container) (lib.Container, error) {
+func (m *maxmindMMDBIn) Input(container lib.Container) (lib.Container, error) {
 	var content []byte
 	var err error
 	switch {
-	case strings.HasPrefix(strings.ToLower(g.URI), "http://"), strings.HasPrefix(strings.ToLower(g.URI), "https://"):
-		content, err = lib.GetRemoteURLContent(g.URI)
+	case strings.HasPrefix(strings.ToLower(m.URI), "http://"), strings.HasPrefix(strings.ToLower(m.URI), "https://"):
+		content, err = lib.GetRemoteURLContent(m.URI)
 	default:
-		content, err = os.ReadFile(g.URI)
+		content, err = os.ReadFile(m.URI)
 	}
 	if err != nil {
 		return nil, err
 	}
 
 	entries := make(map[string]*lib.Entry, 300)
-	err = g.generateEntries(content, entries)
+	err = m.generateEntries(content, entries)
 	if err != nil {
 		return nil, err
 	}
 
 	if len(entries) == 0 {
-		return nil, fmt.Errorf("❌ [type %s | action %s] no entry is generated", typeMaxmindMMDBIn, g.Action)
+		return nil, fmt.Errorf("❌ [type %s | action %s] no entry is generated", typeMaxmindMMDBIn, m.Action)
 	}
 
 	var ignoreIPType lib.IgnoreIPOption
-	switch g.OnlyIPType {
+	switch m.OnlyIPType {
 	case lib.IPv4:
 		ignoreIPType = lib.IgnoreIPv6
 	case lib.IPv6:
 		ignoreIPType = lib.IgnoreIPv4
 	}
 
-	// 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 _, entry := range entries {
-		name := entry.GetName()
-		if len(wantList) > 0 && !wantList[name] {
-			continue
-		}
-
-		switch g.Action {
+		switch m.Action {
 		case lib.ActionAdd:
 			if err := container.Add(entry, ignoreIPType); err != nil {
 				return nil, err
@@ -139,7 +134,7 @@ func (g *maxmindMMDBIn) Input(container lib.Container) (lib.Container, error) {
 	return container, nil
 }
 
-func (g *maxmindMMDBIn) generateEntries(content []byte, entries map[string]*lib.Entry) error {
+func (m *maxmindMMDBIn) generateEntries(content []byte, entries map[string]*lib.Entry) error {
 	db, err := maxminddb.FromBytes(content)
 	if err != nil {
 		return err
@@ -177,6 +172,10 @@ func (g *maxmindMMDBIn) generateEntries(content []byte, entries map[string]*lib.
 			continue
 		}
 
+		if len(m.Want) > 0 && !m.Want[name] {
+			continue
+		}
+
 		entry, found := entries[name]
 		if !found {
 			entry = lib.NewEntry(name)
diff --git a/plugin/maxmind/mmdb_out.go b/plugin/maxmind/mmdb_out.go
index 5b838463..e42ea02d 100644
--- a/plugin/maxmind/mmdb_out.go
+++ b/plugin/maxmind/mmdb_out.go
@@ -140,7 +140,7 @@ func (m *mmdbOut) getEntryNameListInOrder(container lib.Container) []string {
 		The order of names in wantedList has a higher priority than which of the overwriteList.
 	*/
 
-	wantList := make([]string, 0, 200)
+	wantList := make([]string, 0, len(m.Want))
 	for _, want := range m.Want {
 		if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
 			wantList = append(wantList, want)
@@ -151,7 +151,7 @@ func (m *mmdbOut) getEntryNameListInOrder(container lib.Container) []string {
 		return wantList
 	}
 
-	overwriteList := make([]string, 0, 200)
+	overwriteList := make([]string, 0, len(m.Overwrite))
 	overwriteMap := make(map[string]bool)
 	for _, overwrite := range m.Overwrite {
 		if overwrite = strings.ToUpper(strings.TrimSpace(overwrite)); overwrite != "" {