From 84a8d4e64507b0fba5f72600f9c216e2ad7e7008 Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Thu, 9 Jun 2022 19:12:10 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96nftables=E7=9B=B8=E5=85=B3?=
 =?UTF-8?q?=E4=BB=A3=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/firewalls/ddos_protection.go    | 10 +++++++++-
 internal/firewalls/firewall_firewalld.go |  3 ++-
 internal/firewalls/firewall_nftables.go  | 24 +++++++++++++++++++++++-
 internal/nodes/listener_manager.go       |  6 ++++++
 4 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/internal/firewalls/ddos_protection.go b/internal/firewalls/ddos_protection.go
index dd1e5c1..8486fd9 100644
--- a/internal/firewalls/ddos_protection.go
+++ b/internal/firewalls/ddos_protection.go
@@ -17,6 +17,7 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/zero"
 	"github.com/iwind/TeaGo/lists"
 	"github.com/iwind/TeaGo/types"
+	stringutil "github.com/iwind/TeaGo/utils/string"
 	"net"
 	"os/exec"
 	"strings"
@@ -153,6 +154,11 @@ func (this *DDoSProtectionManager) Apply(config *ddosconfigs.ProtectionConfig) e
 
 // 添加TCP规则
 func (this *DDoSProtectionManager) addTCPRules(tcpConfig *ddosconfigs.TCPConfig) error {
+	// 检查nft版本不能小于0.9
+	if len(nftablesInstance.version) > 0 && stringutil.VersionCompare("0.9", nftablesInstance.version) > 0 {
+		return nil
+	}
+
 	var ports = []int32{}
 	for _, portConfig := range tcpConfig.Ports {
 		if !lists.ContainsInt32(ports, portConfig.Port) {
@@ -237,9 +243,11 @@ func (this *DDoSProtectionManager) addTCPRules(tcpConfig *ddosconfigs.TCPConfig)
 		for _, port := range ports {
 			if maxConnections > 0 {
 				var cmd = exec.Command(this.nftPath, "add", "rule", protocol, filter.Name, nftablesChainName, "tcp", "dport", types.String(port), "ct", "count", "over", types.String(maxConnections), "counter", "drop", "comment", this.encodeUserData([]string{"tcp", types.String(port), "maxConnections", types.String(maxConnections)}))
+				var stderr = &bytes.Buffer{}
+				cmd.Stderr = stderr
 				err := cmd.Run()
 				if err != nil {
-					return errors.New("add nftables rule '" + cmd.String() + "' failed: " + err.Error())
+					return errors.New("add nftables rule '" + cmd.String() + "' failed: " + err.Error() + " (" + stderr.String() + ")")
 				}
 			}
 
diff --git a/internal/firewalls/firewall_firewalld.go b/internal/firewalls/firewall_firewalld.go
index 3d63876..a25b304 100644
--- a/internal/firewalls/firewall_firewalld.go
+++ b/internal/firewalls/firewall_firewalld.go
@@ -23,12 +23,13 @@ func NewFirewalld() *Firewalld {
 
 	path, err := exec.LookPath("firewall-cmd")
 	if err == nil && len(path) > 0 {
-		var cmd = exec.Command(path, "-V")
+		var cmd = exec.Command(path, "--state")
 		err := cmd.Run()
 		if err == nil {
 			firewalld.exe = path
 			// TODO check firewalld status with 'firewall-cmd --state' (running or not running),
 			//      but we should recover the state when firewalld state changes, maybe check it every minutes
+
 			firewalld.isReady = true
 			firewalld.init()
 		}
diff --git a/internal/firewalls/firewall_nftables.go b/internal/firewalls/firewall_nftables.go
index 3a959eb..20c2ce7 100644
--- a/internal/firewalls/firewall_nftables.go
+++ b/internal/firewalls/firewall_nftables.go
@@ -5,6 +5,7 @@
 package firewalls
 
 import (
+	"bytes"
 	"errors"
 	"github.com/TeaOSLab/EdgeNode/internal/events"
 	"github.com/TeaOSLab/EdgeNode/internal/firewalls/nftables"
@@ -12,6 +13,7 @@ import (
 	"github.com/iwind/TeaGo/types"
 	"net"
 	"os/exec"
+	"regexp"
 	"runtime"
 	"strings"
 	"time"
@@ -87,6 +89,7 @@ func NewNFTablesFirewall() (*NFTablesFirewall, error) {
 type NFTablesFirewall struct {
 	conn    *nftables.Conn
 	isReady bool
+	version string
 
 	allowIPv4Set *nftables.Set
 	allowIPv6Set *nftables.Set
@@ -99,10 +102,11 @@ type NFTablesFirewall struct {
 
 func (this *NFTablesFirewall) init() error {
 	// check nft
-	_, err := exec.LookPath("nft")
+	nftPath, err := exec.LookPath("nft")
 	if err != nil {
 		return errors.New("nft not found")
 	}
+	this.version = this.readVersion(nftPath)
 
 	// table
 	for _, tableDef := range nftablesFilters {
@@ -371,3 +375,21 @@ func (this *NFTablesFirewall) RemoveSourceIP(ip string) error {
 
 	return nil
 }
+
+// 读取版本号
+func (this *NFTablesFirewall) readVersion(nftPath string) string {
+	var cmd = exec.Command(nftPath, "--version")
+	var output = &bytes.Buffer{}
+	cmd.Stdout = output
+	err := cmd.Run()
+	if err != nil {
+		return ""
+	}
+
+	var outputString = output.String()
+	var versionMatches = regexp.MustCompile(`nftables v([\d.]+)`).FindStringSubmatch(outputString)
+	if len(versionMatches) <= 1 {
+		return ""
+	}
+	return versionMatches[1]
+}
diff --git a/internal/nodes/listener_manager.go b/internal/nodes/listener_manager.go
index d8771be..7ed8771 100644
--- a/internal/nodes/listener_manager.go
+++ b/internal/nodes/listener_manager.go
@@ -257,6 +257,12 @@ func (this *ListenerManager) addToFirewalld(groupAddrs []string) {
 		return
 	}
 
+	// 检查状态
+	err = exec.Command(firewallCmd, "--state").Run()
+	if err != nil {
+		return
+	}
+
 	remotelogs.Println("FIREWALLD", "open ports automatically")
 	for _, port := range ports {
 		{