Fix: remove entry & prefix

2 files changed, 21 insertions, 21 deletions

diff --git a/plugin/maxmind/country_csv.go b/plugin/maxmind/country_csv.go
index 4608f121..aea07a73 100644
--- a/plugin/maxmind/country_csv.go
+++ b/plugin/maxmind/country_csv.go
@@ -4,6 +4,7 @@ import (
 	"encoding/csv"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"os"
 	"path/filepath"
 	"strings"
@@ -113,6 +114,10 @@ func (g *geoLite2CountryCSV) Input(container lib.Container) (lib.Container, erro
 		}
 	}
 
+	if len(entries) == 0 {
+		return nil, fmt.Errorf("❌ [type %s | action %s] no entry is generated", typeCountryCSV, g.Action)
+	}
+
 	var ignoreIPType lib.IgnoreIPOption
 	switch g.OnlyIPType {
 	case lib.IPv4:
@@ -121,14 +126,16 @@ func (g *geoLite2CountryCSV) Input(container lib.Container) (lib.Container, erro
 		ignoreIPType = lib.IgnoreIPv4
 	}
 
-	for name, entry := range entries {
+	for _, entry := range entries {
 		switch g.Action {
 		case lib.ActionAdd:
 			if err := container.Add(entry, ignoreIPType); err != nil {
 				return nil, err
 			}
 		case lib.ActionRemove:
-			container.Remove(name, ignoreIPType)
+			if err := container.Remove(entry, lib.CaseRemovePrefix, ignoreIPType); err != nil {
+				return nil, err
+			}
 		default:
 			return nil, lib.ErrUnknownAction
 		}
@@ -193,10 +200,8 @@ func (g *geoLite2CountryCSV) process(file string, ccMap map[string]string, entri
 	for _, line := range lines[1:] {
 		ccID := strings.TrimSpace(line[1])
 		if countryCode, found := ccMap[ccID]; found {
-			if len(wantList) > 0 {
-				if _, found := wantList[countryCode]; !found {
-					continue
-				}
+			if len(wantList) > 0 && !wantList[countryCode] {
+				continue
 			}
 			cidrStr := strings.ToLower(strings.TrimSpace(line[0]))
 			entry, found := entries[countryCode]
diff --git a/plugin/maxmind/mmdb_in.go b/plugin/maxmind/mmdb_in.go
index eef4c7d6..0c94a12c 100644
--- a/plugin/maxmind/mmdb_in.go
+++ b/plugin/maxmind/mmdb_in.go
@@ -107,7 +107,7 @@ func (g *maxmindMMDBIn) Input(container lib.Container) (lib.Container, error) {
 	}
 
 	if len(entries) == 0 {
-		return nil, fmt.Errorf("❌ [type %s | action %s] no entry is newly generated", typeMaxmindMMDBIn, g.Action)
+		return nil, fmt.Errorf("❌ [type %s | action %s] no entry is generated", typeMaxmindMMDBIn, g.Action)
 	}
 
 	var ignoreIPType lib.IgnoreIPOption
@@ -138,7 +138,11 @@ func (g *maxmindMMDBIn) Input(container lib.Container) (lib.Container, error) {
 				return nil, err
 			}
 		case lib.ActionRemove:
-			container.Remove(name, ignoreIPType)
+			if err := container.Remove(entry, lib.CaseRemovePrefix, ignoreIPType); err != nil {
+				return nil, err
+			}
+		default:
+			return nil, lib.ErrUnknownAction
 		}
 	}
 
@@ -197,23 +201,14 @@ func (g *maxmindMMDBIn) generateEntries(entries map[string]*lib.Entry) error {
 			continue
 		}
 
-		var entry *lib.Entry
 		name := strings.ToUpper(record.Country.IsoCode)
-		if theEntry, found := entries[name]; found {
-			entry = theEntry
-		} else {
+		entry, found := entries[name]
+		if !found {
 			entry = lib.NewEntry(name)
 		}
 
-		switch g.Action {
-		case lib.ActionAdd:
-			if err := entry.AddPrefix(subnet); err != nil {
-				return err
-			}
-		case lib.ActionRemove:
-			if err := entry.RemovePrefix(subnet.String()); err != nil {
-				return err
-			}
+		if err := entry.AddPrefix(subnet); err != nil {
+			return err
 		}
 
 		entries[name] = entry