Refine: wantedList in various formats

2 files changed, 28 insertions, 28 deletions

diff --git a/plugin/v2ray/dat_in.go b/plugin/v2ray/dat_in.go
index 5baf32f2..4b2a98c0 100644
--- a/plugin/v2ray/dat_in.go
+++ b/plugin/v2ray/dat_in.go
@@ -45,12 +45,20 @@ func newGeoIPDatIn(action lib.Action, data json.RawMessage) (lib.InputConverter,
 		return nil, fmt.Errorf("[type %s | action %s] uri must be specified in config", typeGeoIPdatIn, action)
 	}
 
+	// 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 &geoIPDatIn{
 		Type:        typeGeoIPdatIn,
 		Action:      action,
 		Description: descGeoIPdatIn,
 		URI:         tmp.URI,
-		Want:        tmp.Want,
+		Want:        wantList,
 		OnlyIPType:  tmp.OnlyIPType,
 	}, nil
 }
@@ -60,7 +68,7 @@ type geoIPDatIn struct {
 	Action      lib.Action
 	Description string
 	URI         string
-	Want        []string
+	Want        map[string]bool
 	OnlyIPType  lib.IPType
 }
 
@@ -103,20 +111,7 @@ func (g *geoIPDatIn) Input(container lib.Container) (lib.Container, error) {
 		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 {
 		case lib.ActionAdd:
 			if err := container.Add(entry, ignoreIPType); err != nil {
@@ -178,7 +173,12 @@ func (g *geoIPDatIn) generateEntries(reader io.Reader, entries map[string]*lib.E
 	}
 
 	for _, geoip := range geoipList.Entry {
-		name := geoip.CountryCode
+		name := strings.ToUpper(strings.TrimSpace(geoip.CountryCode))
+
+		if len(g.Want) > 0 && !g.Want[name] {
+			continue
+		}
+
 		entry, found := entries[name]
 		if !found {
 			entry = lib.NewEntry(name)
diff --git a/plugin/v2ray/dat_out.go b/plugin/v2ray/dat_out.go
index 0dd9e23b..abb00ca4 100644
--- a/plugin/v2ray/dat_out.go
+++ b/plugin/v2ray/dat_out.go
@@ -57,13 +57,21 @@ func newGeoIPDat(action lib.Action, data json.RawMessage) (lib.OutputConverter,
 		tmp.OutputDir = defaultOutputDir
 	}
 
+	// Filter want list
+	wantList := make([]string, 0, len(tmp.Want))
+	for _, want := range tmp.Want {
+		if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
+			wantList = append(wantList, want)
+		}
+	}
+
 	return &geoIPDatOut{
 		Type:           typeGeoIPdatOut,
 		Action:         action,
 		Description:    descGeoIPdatOut,
 		OutputName:     tmp.OutputName,
 		OutputDir:      tmp.OutputDir,
-		Want:           tmp.Want,
+		Want:           wantList,
 		OneFilePerList: tmp.OneFilePerList,
 		OnlyIPType:     tmp.OnlyIPType,
 	}, nil
@@ -93,19 +101,11 @@ func (g *geoIPDatOut) GetDescription() string {
 }
 
 func (g *geoIPDatOut) Output(container lib.Container) error {
-	// Filter want list
-	wantList := make([]string, 0, 50)
-	for _, want := range g.Want {
-		if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
-			wantList = append(wantList, want)
-		}
-	}
-
 	geoIPList := new(router.GeoIPList)
 	geoIPList.Entry = make([]*router.GeoIP, 0, 300)
 	updated := false
 
-	switch len(wantList) {
+	switch len(g.Want) {
 	case 0:
 		list := make([]string, 0, 300)
 		for entry := range container.Loop() {
@@ -143,9 +143,9 @@ func (g *geoIPDatOut) Output(container lib.Container) error {
 
 	default:
 		// Sort the list
-		sort.Strings(wantList)
+		sort.Strings(g.Want)
 
-		for _, name := range wantList {
+		for _, name := range g.Want {
 			entry, found := container.GetEntry(name)
 			if !found {
 				log.Printf("❌ entry %s not found", name)