nftables封禁IP使用异步操作

cmd/edge-node/main.go

@@ -161,7 +161,7 @@ func main() {
 	app.On("ip.drop", func() {
 		var args = os.Args[2:]
 		if len(args) == 0 {
-			fmt.Println("Usage: edge-node ip.drop IP [--timeout=SECONDS]")
+			fmt.Println("Usage: edge-node ip.drop IP [--timeout=SECONDS] [--async]")
 			return
 		}
 		var ip = args[0]
@@ -175,6 +175,11 @@ func main() {
 		if ok {
 			timeoutSeconds = types.Int(timeout[0])
 		}
+		var async = false
+		_, ok = options["async"]
+		if ok {
+			async = true
+		}
 
 		fmt.Println("drop ip '" + ip + "' for '" + types.String(timeoutSeconds) + "' seconds")
 		var sock = gosock.NewTmpSock(teaconst.ProcessName)
@@ -183,6 +188,7 @@ func main() {
 			Params: map[string]interface{}{
 				"ip":             ip,
 				"timeoutSeconds": timeoutSeconds,
+				"async":          async,
 			},
 		})
 		if err != nil {

internal/firewalls/firewall_firewalld.go

@@ -3,6 +3,7 @@
 package firewalls
 
 import (
+	"errors"
 	"github.com/TeaOSLab/EdgeNode/internal/goman"
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
 	"github.com/iwind/TeaGo/types"
@@ -143,7 +144,7 @@ func (this *Firewalld) RejectSourceIP(ip string, timeoutSeconds int) error {
 	return nil
 }
 
-func (this *Firewalld) DropSourceIP(ip string, timeoutSeconds int) error {
+func (this *Firewalld) DropSourceIP(ip string, timeoutSeconds int, async bool) error {
 	if !this.isReady {
 		return nil
 	}
@@ -156,7 +157,15 @@ func (this *Firewalld) DropSourceIP(ip string, timeoutSeconds int) error {
 		args = append(args, "--timeout="+types.String(timeoutSeconds)+"s")
 	}
 	var cmd = exec.Command(this.exe, args...)
-	this.pushCmd(cmd)
+	if async {
+		this.pushCmd(cmd)
+		return nil
+	}
+
+	err := cmd.Run()
+	if err != nil {
+		return errors.New("run command failed '" + cmd.String() + "': " + err.Error())
+	}
 	return nil
 }
 

internal/firewalls/firewall_interface.go

@@ -23,7 +23,10 @@ type FirewallInterface interface {
 	RejectSourceIP(ip string, timeoutSeconds int) error
 
 	// DropSourceIP 丢弃某个源IP数据
-	DropSourceIP(ip string, timeoutSeconds int) error
+	// ip 要封禁的IP
+	// timeoutSeconds 过期时间
+	// async 是否异步
+	DropSourceIP(ip string, timeoutSeconds int, async bool) error
 
 	// RemoveSourceIP 删除某个源IP
 	RemoveSourceIP(ip string) error

internal/firewalls/firewall_mock.go

@@ -47,7 +47,7 @@ func (this *MockFirewall) RejectSourceIP(ip string, timeoutSeconds int) error {
 }
 
 // DropSourceIP 丢弃某个源IP数据
-func (this *MockFirewall) DropSourceIP(ip string, timeoutSeconds int) error {
+func (this *MockFirewall) DropSourceIP(ip string, timeoutSeconds int, async bool) error {
 	_ = ip
 	_ = timeoutSeconds
 	return nil

internal/firewalls/firewall_nftables.go

@@ -10,6 +10,7 @@ import (
 	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
 	"github.com/TeaOSLab/EdgeNode/internal/events"
 	"github.com/TeaOSLab/EdgeNode/internal/firewalls/nftables"
+	"github.com/TeaOSLab/EdgeNode/internal/goman"
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
 	"github.com/iwind/TeaGo/types"
 	"net"
@@ -28,7 +29,7 @@ func init() {
 
 	if runtime.GOOS == "linux" {
 		var ticker = time.NewTicker(3 * time.Minute)
-		go func() {
+		goman.New(func() {
 			for range ticker.C {
 				// if already ready, we break
 				if nftablesIsReady {
@@ -53,7 +54,7 @@ func init() {
 					break
 				}
 			}
-		}()
+		})
 	}
 }
 
@@ -79,9 +80,16 @@ func (this *nftablesTableDefinition) protocol() string {
 	return "ip"
 }
 
+type blockIPItem struct {
+	action         string
+	ip             string
+	timeoutSeconds int
+}
+
 func NewNFTablesFirewall() (*NFTablesFirewall, error) {
 	var firewall = &NFTablesFirewall{
-		conn: nftables.NewConn(),
+		conn:        nftables.NewConn(),
+		dropIPQueue: make(chan *blockIPItem, 4096),
 	}
 	err := firewall.init()
 	if err != nil {
@@ -103,6 +111,8 @@ type NFTablesFirewall struct {
 	denyIPv6Set *nftables.Set
 
 	firewalld *Firewalld
+
+	dropIPQueue chan *blockIPItem
 }
 
 func (this *NFTablesFirewall) init() error {
@@ -248,6 +258,18 @@ func (this *NFTablesFirewall) init() error {
 	nftablesIsReady = true
 	nftablesInstance = this
 
+	goman.New(func() {
+		for ipItem := range this.dropIPQueue {
+			switch ipItem.action {
+			case "drop":
+				err = this.DropSourceIP(ipItem.ip, ipItem.timeoutSeconds, false)
+				if err != nil {
+					remotelogs.Warn("NFTABLES", "drop ip '"+ipItem.ip+"' failed: "+err.Error())
+				}
+			}
+		}
+	})
+
 	// load firewalld
 	var firewalld = NewFirewalld()
 	if firewalld.IsReady() {
@@ -312,16 +334,29 @@ func (this *NFTablesFirewall) AllowSourceIP(ip string) error {
 // RejectSourceIP 拒绝某个源IP连接
 // we did not create set for drop ip, so we reuse DropSourceIP() method here
 func (this *NFTablesFirewall) RejectSourceIP(ip string, timeoutSeconds int) error {
-	return this.DropSourceIP(ip, timeoutSeconds)
+	return this.DropSourceIP(ip, timeoutSeconds, true)
 }
 
 // DropSourceIP 丢弃某个源IP数据
-func (this *NFTablesFirewall) DropSourceIP(ip string, timeoutSeconds int) error {
+func (this *NFTablesFirewall) DropSourceIP(ip string, timeoutSeconds int, async bool) error {
 	var data = net.ParseIP(ip)
 	if data == nil {
 		return errors.New("invalid ip '" + ip + "'")
 	}
 
+	if async {
+		select {
+		case this.dropIPQueue <- &blockIPItem{
+			action:         "drop",
+			ip:             ip,
+			timeoutSeconds: timeoutSeconds,
+		}:
+		default:
+			return errors.New("drop ip queue is full")
+		}
+		return nil
+	}
+
 	if strings.Contains(ip, ":") { // ipv6
 		if this.denyIPv6Set == nil {
 			return errors.New("ipv6 ip set is nil")

internal/firewalls/firewall_nftables_others.go

@@ -51,7 +51,7 @@ func (this *NFTablesFirewall) RejectSourceIP(ip string, timeoutSeconds int) erro
 }
 
 // DropSourceIP 丢弃某个源IP数据
-func (this *NFTablesFirewall) DropSourceIP(ip string, timeoutSeconds int) error {
+func (this *NFTablesFirewall) DropSourceIP(ip string, timeoutSeconds int, async bool) error {
 	return nil
 }
 

internal/nodes/client_listener.go

@@ -57,7 +57,7 @@ func (this *ClientListener) Accept() (net.Conn, error) {
 			if beingDenied {
 				var fw = firewalls.Firewall()
 				if fw != nil && !fw.IsMock() {
-					_ = fw.DropSourceIP(ip, 60)
+					_ = fw.DropSourceIP(ip, 120, true)
 				}
 			}
 

internal/nodes/node.go

@@ -784,7 +784,8 @@ func (this *Node) listenSock() error {
 				var m = maps.NewMap(cmd.Params)
 				var ip = m.GetString("ip")
 				var timeSeconds = m.GetInt("timeoutSeconds")
-				err := firewalls.Firewall().DropSourceIP(ip, timeSeconds)
+				var async = m.GetBool("async")
+				err := firewalls.Firewall().DropSourceIP(ip, timeSeconds, async)
 				if err != nil {
 					_ = cmd.Reply(&gosock.Command{
 						Params: map[string]interface{}{

internal/waf/ip_list.go

@@ -112,7 +112,7 @@ func (this *IPList) RecordIP(ipType string,
 				if seconds > 3600 {
 					seconds = 3600
 				}
-				_ = firewalls.Firewall().DropSourceIP(ip, int(seconds))
+				_ = firewalls.Firewall().DropSourceIP(ip, int(seconds), true)
 			}
 		}
 	}