From 4cf50641990c7df79b1a03438dd9273532182e2f Mon Sep 17 00:00:00 2001
From: Loyalsoldier <10487845+Loyalsoldier@users.noreply.github.com>
Date: Tue, 22 Oct 2024 05:15:50 +0800
Subject: Feat: support to specify `onlyIPType` option in private input format

---
 plugin/special/private.go | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

(limited to 'plugin')

diff --git a/plugin/special/private.go b/plugin/special/private.go
index 4250c225..dc781c31 100644
--- a/plugin/special/private.go
+++ b/plugin/special/private.go
@@ -46,10 +46,21 @@ func init() {
 }
 
 func newPrivate(action lib.Action, data json.RawMessage) (lib.InputConverter, error) {
+	var tmp struct {
+		OnlyIPType lib.IPType `json:"onlyIPType"`
+	}
+
+	if len(data) > 0 {
+		if err := json.Unmarshal(data, &tmp); err != nil {
+			return nil, err
+		}
+	}
+
 	return &private{
 		Type:        typePrivate,
 		Action:      action,
 		Description: descPrivate,
+		OnlyIPType:  tmp.OnlyIPType,
 	}, nil
 }
 
@@ -57,6 +68,7 @@ type private struct {
 	Type        string
 	Action      lib.Action
 	Description string
+	OnlyIPType  lib.IPType
 }
 
 func (p *private) GetType() string {
@@ -83,13 +95,21 @@ func (p *private) Input(container lib.Container) (lib.Container, error) {
 		}
 	}
 
+	var ignoreIPType lib.IgnoreIPOption
+	switch p.OnlyIPType {
+	case lib.IPv4:
+		ignoreIPType = lib.IgnoreIPv6
+	case lib.IPv6:
+		ignoreIPType = lib.IgnoreIPv4
+	}
+
 	switch p.Action {
 	case lib.ActionAdd:
-		if err := container.Add(entry); err != nil {
+		if err := container.Add(entry, ignoreIPType); err != nil {
 			return nil, err
 		}
 	case lib.ActionRemove:
-		if err := container.Remove(entry, lib.CaseRemovePrefix); err != nil {
+		if err := container.Remove(entry, lib.CaseRemovePrefix, ignoreIPType); err != nil {
 			return nil, err
 		}
 	default:
-- 
cgit v1.3.1