Refactor all plugins to use functional options patterncopilot/refactor-plugins-functional-options

Agent-Logs-Url: https://github.com/Loyalsoldier/geoip/sessions/e2b66c9a-3d01-490c-9b31-32109cfe4feb

Co-authored-by: Loyalsoldier <[email protected]>

1 files changed, 95 insertions, 58 deletions

diff --git a/plugin/maxmind/maxmind_asn_csv_in.go b/plugin/maxmind/maxmind_asn_csv_in.go
index d924e3f9..8444277a 100644
--- a/plugin/maxmind/maxmind_asn_csv_in.go
+++ b/plugin/maxmind/maxmind_asn_csv_in.go
@@ -24,101 +24,138 @@ var (
 
 func init() {
 	lib.RegisterInputConfigCreator(TypeGeoLite2ASNCSVIn, func(action lib.Action, data json.RawMessage) (lib.InputConverter, error) {
-		return newGeoLite2ASNCSVIn(action, data)
+		return NewGeoLite2ASNCSVInFromBytes(action, data)
 	})
-	lib.RegisterInputConverter(TypeGeoLite2ASNCSVIn, &GeoLite2ASNCSVIn{
+	lib.RegisterInputConverter(TypeGeoLite2ASNCSVIn, &geoLite2ASNCSVIn{
 		Description: DescGeoLite2ASNCSVIn,
 	})
 }
 
-func newGeoLite2ASNCSVIn(action lib.Action, data json.RawMessage) (lib.InputConverter, error) {
-	var tmp struct {
-		IPv4File   string                 `json:"ipv4"`
-		IPv6File   string                 `json:"ipv6"`
-		Want       lib.WantedListExtended `json:"wantedList"`
-		OnlyIPType lib.IPType             `json:"onlyIPType"`
+type geoLite2ASNCSVIn struct {
+	Type        string
+	Action      lib.Action
+	Description string
+	IPv4File    string
+	IPv6File    string
+	Want        map[string][]string
+	OnlyIPType  lib.IPType
+}
+
+func NewGeoLite2ASNCSVIn(action lib.Action, opts ...lib.InputOption) lib.InputConverter {
+	g := &geoLite2ASNCSVIn{
+		Type:        TypeGeoLite2ASNCSVIn,
+		Action:      action,
+		Description: DescGeoLite2ASNCSVIn,
 	}
 
-	if len(data) > 0 {
-		if err := json.Unmarshal(data, &tmp); err != nil {
-			return nil, err
+	for _, opt := range opts {
+		if opt != nil {
+			opt(g)
 		}
 	}
 
-	// When both of IP files are not specified,
-	// it means user wants to use the default ones
-	if tmp.IPv4File == "" && tmp.IPv6File == "" {
-		tmp.IPv4File = defaultGeoLite2ASNCSVIPv4File
-		tmp.IPv6File = defaultGeoLite2ASNCSVIPv6File
+	return g
+}
+
+func WithGeoLite2ASNCSVInIPv4File(file string) lib.InputOption {
+	return func(s lib.InputConverter) {
+		s.(*geoLite2ASNCSVIn).IPv4File = file
 	}
+}
+
+func WithGeoLite2ASNCSVInIPv6File(file string) lib.InputOption {
+	return func(s lib.InputConverter) {
+		s.(*geoLite2ASNCSVIn).IPv6File = file
+	}
+}
+
+func WithGeoLite2ASNCSVInWantedList(want lib.WantedListExtended) lib.InputOption {
+	return func(s lib.InputConverter) {
+		wantList := make(map[string][]string)
 
-	// Filter want list
-	wantList := make(map[string][]string) // map[asn][]listname or map[asn][]asn
+		for list, asnList := range want.TypeMap {
+			list = strings.ToUpper(strings.TrimSpace(list))
+			if list == "" {
+				continue
+			}
+
+			for _, asn := range asnList {
+				asn = strings.TrimPrefix(strings.ToLower(strings.TrimSpace(asn)), "as")
+				if asn == "" {
+					continue
+				}
 
-	for list, asnList := range tmp.Want.TypeMap {
-		list = strings.ToUpper(strings.TrimSpace(list))
-		if list == "" {
-			continue
+				if listArr, found := wantList[asn]; found {
+					listArr = append(listArr, list)
+					wantList[asn] = listArr
+				} else {
+					wantList[asn] = []string{list}
+				}
+			}
 		}
 
-		for _, asn := range asnList {
+		for _, asn := range want.TypeSlice {
 			asn = strings.TrimPrefix(strings.ToLower(strings.TrimSpace(asn)), "as")
 			if asn == "" {
 				continue
 			}
 
-			if listArr, found := wantList[asn]; found {
-				listArr = append(listArr, list)
-				wantList[asn] = listArr
-			} else {
-				wantList[asn] = []string{list}
-			}
+			wantList[asn] = []string{"AS" + asn}
 		}
+
+		s.(*geoLite2ASNCSVIn).Want = wantList
 	}
+}
 
-	for _, asn := range tmp.Want.TypeSlice {
-		asn = strings.TrimPrefix(strings.ToLower(strings.TrimSpace(asn)), "as")
-		if asn == "" {
-			continue
-		}
+func WithGeoLite2ASNCSVInOnlyIPType(onlyIPType lib.IPType) lib.InputOption {
+	return func(s lib.InputConverter) {
+		s.(*geoLite2ASNCSVIn).OnlyIPType = onlyIPType
+	}
+}
 
-		wantList[asn] = []string{"AS" + asn}
+func NewGeoLite2ASNCSVInFromBytes(action lib.Action, data []byte) (lib.InputConverter, error) {
+	var tmp struct {
+		IPv4File   string                 `json:"ipv4"`
+		IPv6File   string                 `json:"ipv6"`
+		Want       lib.WantedListExtended `json:"wantedList"`
+		OnlyIPType lib.IPType             `json:"onlyIPType"`
 	}
 
-	return &GeoLite2ASNCSVIn{
-		Type:        TypeGeoLite2ASNCSVIn,
-		Action:      action,
-		Description: DescGeoLite2ASNCSVIn,
-		IPv4File:    tmp.IPv4File,
-		IPv6File:    tmp.IPv6File,
-		Want:        wantList,
-		OnlyIPType:  tmp.OnlyIPType,
-	}, nil
-}
+	if len(data) > 0 {
+		if err := json.Unmarshal(data, &tmp); err != nil {
+			return nil, err
+		}
+	}
 
-type GeoLite2ASNCSVIn struct {
-	Type        string
-	Action      lib.Action
-	Description string
-	IPv4File    string
-	IPv6File    string
-	Want        map[string][]string
-	OnlyIPType  lib.IPType
+	// When both of IP files are not specified,
+	// it means user wants to use the default ones
+	if tmp.IPv4File == "" && tmp.IPv6File == "" {
+		tmp.IPv4File = defaultGeoLite2ASNCSVIPv4File
+		tmp.IPv6File = defaultGeoLite2ASNCSVIPv6File
+	}
+
+	return NewGeoLite2ASNCSVIn(
+		action,
+		WithGeoLite2ASNCSVInIPv4File(tmp.IPv4File),
+		WithGeoLite2ASNCSVInIPv6File(tmp.IPv6File),
+		WithGeoLite2ASNCSVInWantedList(tmp.Want),
+		WithGeoLite2ASNCSVInOnlyIPType(tmp.OnlyIPType),
+	), nil
 }
 
-func (g *GeoLite2ASNCSVIn) GetType() string {
+func (g *geoLite2ASNCSVIn) GetType() string {
 	return g.Type
 }
 
-func (g *GeoLite2ASNCSVIn) GetAction() lib.Action {
+func (g *geoLite2ASNCSVIn) GetAction() lib.Action {
 	return g.Action
 }
 
-func (g *GeoLite2ASNCSVIn) GetDescription() string {
+func (g *geoLite2ASNCSVIn) GetDescription() string {
 	return g.Description
 }
 
-func (g *GeoLite2ASNCSVIn) Input(container lib.Container) (lib.Container, error) {
+func (g *geoLite2ASNCSVIn) Input(container lib.Container) (lib.Container, error) {
 	entries := make(map[string]*lib.Entry)
 
 	if g.IPv4File != "" {
@@ -157,7 +194,7 @@ func (g *GeoLite2ASNCSVIn) Input(container lib.Container) (lib.Container, error)
 	return container, nil
 }
 
-func (g *GeoLite2ASNCSVIn) process(file string, entries map[string]*lib.Entry) error {
+func (g *geoLite2ASNCSVIn) process(file string, entries map[string]*lib.Entry) error {
 	if entries == nil {
 		entries = make(map[string]*lib.Entry)
 	}