新架构

2025-12-23 13:14:53 +08:00
parent 7c783f7bdf
commit 340509fd9f
9 changed files with 464 additions and 494 deletions
--- a/ipset.go
+++ b/ipset.go
@@ -4,181 +4,154 @@ import (
 	"bytes"
 	"fmt"
 	"log"
+	"os"
 	"os/exec"
 	"strconv"
 	"strings"
 )

+// 创建 ipset 集合
 func createIPSet(setName string) error {
-	cmd := exec.Command("ipset", "create", setName, "hash:ip")
-	var stdout, stderr bytes.Buffer
-	cmd.Stdout = &stdout
+	// 建议增加 maxelem，防止频繁扩容。默认 65536 有点小。
+	// 加上 -exist 防止报错
+	cmd := exec.Command("ipset", "create", setName, "hash:ip", "maxelem", "200000", "-exist")
+	var stderr bytes.Buffer
 	cmd.Stderr = &stderr

 	err := cmd.Run()
 	if err != nil {
-		// 记录错误信息，但不退出
-		log.Printf("failed to execute command: %v, stderr: %s", err, stderr.String())
+		return fmt.Errorf("ipset create failed: %v, stderr: %s", err, stderr.String())
 	}
-	return err // 返回错误以便调用者处理
+	return nil
 }

+// 添加 IP 到集合
 func AddIPSet(setName string, ip string) error {
 	cmd := exec.Command("ipset", "add", setName, ip)
-	var stdout, stderr bytes.Buffer
-	cmd.Stdout = &stdout
-	cmd.Stderr = &stderr
-
-	err := cmd.Run()
-	if err != nil {
-		// 记录错误信息，但不退出
-		log.Printf("failed to add IP to set: %v, stderr: %s", err, stderr.String())
-	}
-	return err // 返回错误以便调用者处理
+	// 不关注输出，只关注 exit code
+	return cmd.Run()
 }

-// NumIPSet returns the number of entries in the specified ipset set.
+// NumIPSet 获取集合中 IP 的数量
+// 注意：频繁调用此函数开销较大（因为它通过 shell 解析文本），建议仅在监控循环中低频调用
 func NumIPSet(setName string) (int, error) {
-	cmd := exec.Command("sh", "-c", fmt.Sprintf("ipset list %s | grep \"Number of entries\" | cut -d ':' -f 2 | sed 's/ //g'", setName))
+	// 这里必须用 sh -c 因为涉及管道符 |
+	cmd := exec.Command("sh", "-c", fmt.Sprintf("ipset list %s | grep 'Number of entries' | cut -d ':' -f 2", setName))

-	var stdout, stderr bytes.Buffer
+	var stdout bytes.Buffer
 	cmd.Stdout = &stdout
-	cmd.Stderr = &stderr

 	err := cmd.Run()
 	if err != nil {
-		log.Printf("cmd.Run() failed with %v, stderr: %s\n", err, stderr.String())
-		return 0, fmt.Errorf("failed to execute command: %w, stderr: %s", err, stderr.String())
+		return 0, err
 	}

 	output := strings.TrimSpace(stdout.String())
 	numEntries, err := strconv.Atoi(output)
 	if err != nil {
-		log.Printf("failed to parse output as integer: %v, output: %s\n", err, output)
-		return 0, fmt.Errorf("failed to parse output as integer: %w, output: %s", err, output)
+		return 0, fmt.Errorf("解析数量失败: %w, 输出: %s", err, output)
 	}

 	return numEntries, nil
 }

-// IsIpset 检查名为 setName 的 ipset 是否存在，通过返回 0 表示存在，非零表示不存在或其他错误。
+// Is_Name_Ipset 检查集合是否存在
+// 返回 0 表示存在 (ExitCode 0)，其他表示不存在
 func Is_Name_Ipset(setName string) int {
-	cmd := exec.Command("ipset", "list", setName)
+	cmd := exec.Command("ipset", "list", setName, "-n") // -n 只输出名字，减少开销
 	err := cmd.Run()
-
 	if err != nil {
-		if exitError, ok := err.(*exec.ExitError); ok {
-			// The program has exited with an exit code != 0
-			return exitError.ExitCode()
-		} else {
-			// Another error occurred (e.g., command not found)
-			return -1 // 或者你可以选择其他方式来标识这种情况
-		}
+		return 1 // 不存在或出错
 	}
-
-	// Command executed successfully, the set exists
-	return 0
+	return 0 // 存在
 }

+// Is_Ip_Ipset 检查 IP 是否在 ANY root 集合中
+// 优化：不再使用 grep 全局搜索，而是遍历 root0...rootN 使用 ipset test
 func Is_Ip_Ipset(ip string) int {
-	cmd := exec.Command("sh", "-c", fmt.Sprintf("ipset list | grep \"%s\"", ip))
-	err := cmd.Run()
+	// 遍历所有可能的集合名
+	for i := 0; i < MAX_IPSET_NAME; i++ {
+		setName := fmt.Sprintf("root%d", i)

-	if err != nil {
-		if exitError, ok := err.(*exec.ExitError); ok {
-			// The program has exited with an exit code != 0
-			return exitError.ExitCode()
-		} else {
-			// Another error occurred (e.g., command not found)
-			return -1 // 或者你可以选择其他方式来标识这种情况
+		// 快速检查集合是否存在（可选，为了严谨）
+		// 如果 ipset test 在集合不存在时会报错，所以我们直接运行 test
+		// 如果 ip 在集合中，exit code 为 0
+		cmd := exec.Command("ipset", "test", setName, ip)
+		err := cmd.Run()
+
+		if err == nil {
+			return 1 // 找到了，IP 在该集合中
 		}
 	}

-	// Command executed successfully, the set exists
-	return 0
+	return 0 // 所有集合都找过了，不在其中
 }

-// RemoveIPIfInSets 在多个 ipset 中查找并删除某个 IP，成功删除则返回集合名
+// RemoveIPIfInSets 在所有集合中查找并删除某个 IP
 func RemoveIPIfInSets(prefix string, max int, ip string) (string, error) {
 	for i := 0; i < max; i++ {
 		setName := fmt.Sprintf("%s%d", prefix, i)

-		// 检查 IP 是否在当前集合中
-		cmd := exec.Command("ipset", "test", setName, ip)
-		output, err := cmd.CombinedOutput()
-
-		if err != nil {
-			if strings.Contains(string(output), "is NOT in set") {
-				continue // 不在该集合中，尝试下一个
-			}
-			return "", fmt.Errorf("检测 %s 时出错: %v (%s)", setName, err, output)
+		// 1. 先测试是否存在 (避免无意义的 del 调用报错)
+		testCmd := exec.Command("ipset", "test", setName, ip)
+		if err := testCmd.Run(); err != nil {
+			// exit code != 0 说明不在这个集合里，继续找下一个
+			continue
 		}

-		// 存在，执行删除
-		cmd = exec.Command("ipset", "del", setName, ip)
-		output, err = cmd.CombinedOutput()
-		if err != nil {
-			return "", fmt.Errorf("从 %s 删除 %s 时出错: %v (%s)", setName, ip, err, output)
+		// 2. 存在则删除
+		delCmd := exec.Command("ipset", "del", setName, ip)
+		if err := delCmd.Run(); err != nil {
+			return "", fmt.Errorf("从 %s 删除 %s 失败: %v", setName, ip, err)
 		}

-		return setName, nil // 删除成功，返回集合名
+		return setName, nil // 成功删除
 	}

-	return "", nil // 所有集合都没有该 IP，无需删除
+	return "", nil // 未找到
 }

-// 添加 Iptables 规则
+// iptables_add 添加规则
+// 修正：使用 -I (Insert) 而不是 -A (Append)，确保规则在最前面生效
 func iptables_add(setName string) error {
+	// 检查规则是否已经存在，避免重复添加
+	checkCmd := exec.Command("iptables", "-C", "INPUT", "-m", "set", "--match-set", setName, "src", "-j", "DROP")
+	if err := checkCmd.Run(); err == nil {
+		// 规则已存在，直接返回
+		return nil
+	}

-	cmd := exec.Command("sh", "-c", fmt.Sprintf("iptables -A INPUT -m set --match-set %s src -j DROP", setName))
-
-	var stdout, stderr bytes.Buffer
-	cmd.Stdout = &stdout
+	// 添加规则
+	cmd := exec.Command("iptables", "-I", "INPUT", "-m", "set", "--match-set", setName, "src", "-j", "DROP")
+	var stderr bytes.Buffer
 	cmd.Stderr = &stderr

 	err := cmd.Run()
 	if err != nil {
-		//log.Printf("cmd.Run() failed with %v, stderr: %s\n", err, stderr.String())
-		//err = fmt.Errorf("failed to execute command: %w, stderr: %s", err, stderr.String())
+		log.Printf("添加 iptables 规则失败 [%s]: %v, stderr: %s", setName, err, stderr.String())
 	}
-
 	return err
 }

-// 删除 Iptables 规则
+// iptables_del 删除规则
 func iptables_del(setName string) error {
-
-	cmd := exec.Command("sh", "-c", fmt.Sprintf("iptables -D INPUT -m set --match-set %s src -j DROP", setName))
-
-	var stdout, stderr bytes.Buffer
-	cmd.Stdout = &stdout
-	cmd.Stderr = &stderr
-
-	err := cmd.Run()
-	if err != nil {
-		//log.Printf("cmd.Run() failed with %v, stderr: %s\n", err, stderr.String())
-		//err = fmt.Errorf("failed to execute command: %w, stderr: %s", err, stderr.String())
+	// 循环删除，防止有重复规则
+	for {
+		cmd := exec.Command("iptables", "-D", "INPUT", "-m", "set", "--match-set", setName, "src", "-j", "DROP")
+		if err := cmd.Run(); err != nil {
+			// 删除失败通常意味着规则不存在了，跳出循环
+			break
+		}
 	}
-
-	return err
+	return nil
 }

-// 打印 Iptables 规则
+// iptables_list 打印规则
 func iptables_list() error {
-	cmd := exec.Command("sh", "-c", "iptables -L -v -n --line-numbers")
-
-	var stdout, stderr bytes.Buffer
-	cmd.Stdout = &stdout
-	cmd.Stderr = &stderr
-
-	err := cmd.Run()
-	if err != nil {
-		log.Printf("cmd.Run() failed with %v, stderr: %s\n", err, stderr.String())
-		err = fmt.Errorf("failed to execute command: %w, stderr: %s", err, stderr.String())
-	}
-
-	fmt.Print(stdout.String())
-	return err
+	cmd := exec.Command("iptables", "-L", "INPUT", "-v", "-n", "--line-numbers")
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	return cmd.Run()
 }

 /*