From 6852195a9072358191b6989faaed333831fb1e94 Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Sun, 2 Apr 2023 18:37:24 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96nftables=E5=8F=AF=E6=89=A7?=
 =?UTF-8?q?=E8=A1=8C=E6=96=87=E4=BB=B6=E6=9F=A5=E6=89=BE=E6=96=B9=E6=B3=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/firewalls/ddos_protection.go   | 10 ++------
 internal/firewalls/firewall_nftables.go | 33 +++++++++++++++++++------
 2 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/internal/firewalls/ddos_protection.go b/internal/firewalls/ddos_protection.go
index 9054b62..ecb5c8f 100644
--- a/internal/firewalls/ddos_protection.go
+++ b/internal/firewalls/ddos_protection.go
@@ -20,7 +20,6 @@ import (
 	"github.com/iwind/TeaGo/types"
 	stringutil "github.com/iwind/TeaGo/utils/string"
 	"net"
-	"os/exec"
 	"strings"
 	"time"
 )
@@ -91,7 +90,7 @@ func (this *DDoSProtectionManager) Apply(config *ddosconfigs.ProtectionConfig) e
 	}
 	remotelogs.Println("FIREWALL", "change DDoS protection config")
 
-	if len(this.nftExe()) == 0 {
+	if len(NftExePath()) == 0 {
 		return errors.New("can not find nft command")
 	}
 
@@ -157,7 +156,7 @@ func (this *DDoSProtectionManager) Apply(config *ddosconfigs.ProtectionConfig) e
 
 // 添加TCP规则
 func (this *DDoSProtectionManager) addTCPRules(tcpConfig *ddosconfigs.TCPConfig) error {
-	var nftExe = this.nftExe()
+	var nftExe = NftExePath()
 	if len(nftExe) == 0 {
 		return nil
 	}
@@ -557,8 +556,3 @@ func (this *DDoSProtectionManager) updateAllowIPList(allIPList []string) error {
 
 	return nil
 }
-
-func (this *DDoSProtectionManager) nftExe() string {
-	path, _ := exec.LookPath("nft")
-	return path
-}
diff --git a/internal/firewalls/firewall_nftables.go b/internal/firewalls/firewall_nftables.go
index feeddbe..59a88f5 100644
--- a/internal/firewalls/firewall_nftables.go
+++ b/internal/firewalls/firewall_nftables.go
@@ -1,6 +1,5 @@
 // Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
 //go:build linux
-// +build linux
 
 package firewalls
 
@@ -16,6 +15,7 @@ import (
 	"github.com/google/nftables/expr"
 	"github.com/iwind/TeaGo/types"
 	"net"
+	"os"
 	"os/exec"
 	"regexp"
 	"runtime"
@@ -23,6 +23,25 @@ import (
 	"time"
 )
 
+// NftExePath 查找nftables可执行文件路径
+func NftExePath() string {
+	path, _ := exec.LookPath("nft")
+	if len(path) > 0 {
+		return path
+	}
+
+	for _, possiblePath := range []string{
+		"/usr/sbin/nft",
+	} {
+		_, err := os.Stat(possiblePath)
+		if err == nil {
+			return possiblePath
+		}
+	}
+
+	return ""
+}
+
 // check nft status, if being enabled we load it automatically
 func init() {
 	if !teaconst.IsMain {
@@ -38,8 +57,8 @@ func init() {
 					ticker.Stop()
 					break
 				}
-				_, err := exec.LookPath("nft")
-				if err == nil {
+				var nftExe = NftExePath()
+				if len(nftExe) > 0 {
 					nftablesFirewall, err := NewNFTablesFirewall()
 					if err != nil {
 						continue
@@ -121,9 +140,9 @@ type NFTablesFirewall struct {
 
 func (this *NFTablesFirewall) init() error {
 	// check nft
-	nftPath, err := exec.LookPath("nft")
-	if err != nil {
-		return errors.New("nft not found")
+	var nftPath = NftExePath()
+	if len(nftPath) == 0 {
+		return errors.New("'nft' not found")
 	}
 	this.version = this.readVersion(nftPath)
 
@@ -276,7 +295,7 @@ func (this *NFTablesFirewall) init() error {
 		for ipItem := range this.dropIPQueue {
 			switch ipItem.action {
 			case "drop":
-				err = this.DropSourceIP(ipItem.ip, ipItem.timeoutSeconds, false)
+				err := this.DropSourceIP(ipItem.ip, ipItem.timeoutSeconds, false)
 				if err != nil {
 					remotelogs.Warn("NFTABLES", "drop ip '"+ipItem.ip+"' failed: "+err.Error())
 				}