From f8f0153664984edba115771b72c40fa8a4e8908c Mon Sep 17 00:00:00 2001
From: Toby <tobyxdd@gmail.com>
Date: Sat, 3 Feb 2024 10:55:20 -0800
Subject: [PATCH] feat: rules hot reload via SIGHUP (#44)

---
 README.ja.md     |  1 +
 README.md        |  1 +
 README.zh.md     |  1 +
 cmd/root.go      | 35 ++++++++++++++++++++++++++++++++---
 engine/worker.go |  5 ++++-
 5 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/README.ja.md b/README.ja.md
index 390d262..86fa4a1 100644
--- a/README.ja.md
+++ b/README.ja.md
@@ -25,6 +25,7 @@ OpenGFW は、Linux 上の [GFW](https://en.wikipedia.org/wiki/Great_Firewall) 
 - フローベースのマルチコア負荷分散
 - 接続オフロード
 - [expr](https://github.com/expr-lang/expr) に基づく強力なルールエンジン
+- ルールのホットリロード (`SIGHUP` を送信してリロード)
 - 柔軟なアナライザ＆モディファイアフレームワーク
 - 拡張可能な IO 実装(今のところ NFQueue のみ)
 - [WIP] ウェブ UI
diff --git a/README.md b/README.md
index 55ff4e1..7262463 100644
--- a/README.md
+++ b/README.md
@@ -29,6 +29,7 @@ Linux that's in many ways more powerful than the real thing. It's cyber sovereig
 - Flow-based multicore load balancing
 - Connection offloading
 - Powerful rule engine based on [expr](https://github.com/expr-lang/expr)
+- Hot-reloadable rules (send `SIGHUP` to reload)
 - Flexible analyzer & modifier framework
 - Extensible IO implementation (only NFQueue for now)
 - [WIP] Web UI
diff --git a/README.zh.md b/README.zh.md
index 7a8f2ec..3f3778f 100644
--- a/README.zh.md
+++ b/README.zh.md
@@ -25,6 +25,7 @@ OpenGFW 是一个 Linux 上灵活、易用、开源的 [GFW](https://zh.wikipedi
 - 基于流的多核负载均衡
 - 连接 offloading
 - 基于 [expr](https://github.com/expr-lang/expr) 的强大规则引擎
+- 规则可以热重载 (发送 `SIGHUP` 信号)
 - 灵活的协议解析和修改框架
 - 可扩展的 IO 实现 (目前只有 NFQueue)
 - [开发中] Web UI
diff --git a/cmd/root.go b/cmd/root.go
index f3ff665..7de9159 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -7,6 +7,7 @@ import (
 	"os/signal"
 	"strconv"
 	"strings"
+	"syscall"
 
 	"github.com/apernet/OpenGFW/analyzer"
 	"github.com/apernet/OpenGFW/analyzer/tcp"
@@ -267,14 +268,42 @@ func runMain(cmd *cobra.Command, args []string) {
 		logger.Fatal("failed to initialize engine", zap.Error(err))
 	}
 
+	// Signal handling
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	go func() {
-		sigChan := make(chan os.Signal)
-		signal.Notify(sigChan, os.Interrupt, os.Kill)
-		<-sigChan
+		// Graceful shutdown
+		shutdownChan := make(chan os.Signal)
+		signal.Notify(shutdownChan, os.Interrupt, os.Kill)
+		<-shutdownChan
 		logger.Info("shutting down gracefully...")
 		cancelFunc()
 	}()
+	go func() {
+		// Rule reload
+		reloadChan := make(chan os.Signal)
+		signal.Notify(reloadChan, syscall.SIGHUP)
+		for {
+			<-reloadChan
+			logger.Info("reloading rules")
+			rawRs, err := ruleset.ExprRulesFromYAML(args[0])
+			if err != nil {
+				logger.Error("failed to load rules, using old rules", zap.Error(err))
+				continue
+			}
+			rs, err := ruleset.CompileExprRules(rawRs, analyzers, modifiers, rsConfig)
+			if err != nil {
+				logger.Error("failed to compile rules, using old rules", zap.Error(err))
+				continue
+			}
+			err = en.UpdateRuleset(rs)
+			if err != nil {
+				logger.Error("failed to update ruleset", zap.Error(err))
+			} else {
+				logger.Info("rules reloaded")
+			}
+		}
+	}()
+
 	logger.Info("engine started")
 	logger.Info("engine exited", zap.Error(en.Run(ctx)))
 }
diff --git a/engine/worker.go b/engine/worker.go
index 2bca8e0..5177016 100644
--- a/engine/worker.go
+++ b/engine/worker.go
@@ -127,7 +127,10 @@ func (w *worker) Run(ctx context.Context) {
 }
 
 func (w *worker) UpdateRuleset(r ruleset.Ruleset) error {
-	return w.tcpStreamFactory.UpdateRuleset(r)
+	if err := w.tcpStreamFactory.UpdateRuleset(r); err != nil {
+		return err
+	}
+	return w.udpStreamFactory.UpdateRuleset(r)
 }
 
 func (w *worker) handle(streamID uint32, p gopacket.Packet) (io.Verdict, []byte) {