fix: netlink race condition (#48)

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

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

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

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)))
 }

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) {

io/nfqueue.go

@@ -41,10 +41,11 @@ var _ PacketIO = (*nfqueuePacketIO)(nil)
 var errNotNFQueuePacket = errors.New("not an NFQueue packet")
 
 type nfqueuePacketIO struct {
-	n     *nfqueue.Nfqueue
-	local bool
-	ipt4  *iptables.IPTables
-	ipt6  *iptables.IPTables
+	n      *nfqueue.Nfqueue
+	local  bool
+	ipt4   *iptables.IPTables
+	ipt6   *iptables.IPTables
+	iptSet bool // whether iptables rules are set
 }
 
 type NFQueuePacketIOConfig struct {
@@ -74,22 +75,16 @@ func NewNFQueuePacketIO(config NFQueuePacketIOConfig) (PacketIO, error) {
 	if err != nil {
 		return nil, err
 	}
-	io := &nfqueuePacketIO{
+	return &nfqueuePacketIO{
 		n:     n,
 		local: config.Local,
 		ipt4:  ipt4,
 		ipt6:  ipt6,
-	}
-	err = io.setupIpt(config.Local, false)
-	if err != nil {
-		_ = n.Close()
-		return nil, err
-	}
-	return io, nil
+	}, nil
 }
 
 func (n *nfqueuePacketIO) Register(ctx context.Context, cb PacketCallback) error {
-	return n.n.RegisterWithErrorFunc(ctx,
+	err := n.n.RegisterWithErrorFunc(ctx,
 		func(a nfqueue.Attribute) int {
 			if a.PacketID == nil || a.Ct == nil || a.Payload == nil || len(*a.Payload) < 20 {
 				// Invalid packet, ignore
@@ -106,6 +101,17 @@ func (n *nfqueuePacketIO) Register(ctx context.Context, cb PacketCallback) error
 		func(e error) int {
 			return okBoolToInt(cb(nil, e))
 		})
+	if err != nil {
+		return err
+	}
+	if !n.iptSet {
+		err = n.setupIpt(n.local, false)
+		if err != nil {
+			return err
+		}
+		n.iptSet = true
+	}
+	return nil
 }
 
 func (n *nfqueuePacketIO) SetVerdict(p Packet, v Verdict, newPacket []byte) error {
@@ -150,9 +156,13 @@ func (n *nfqueuePacketIO) setupIpt(local, remove bool) error {
 }
 
 func (n *nfqueuePacketIO) Close() error {
-	err := n.setupIpt(n.local, true)
-	_ = n.n.Close()
-	return err
+	if n.iptSet {
+		err := n.setupIpt(n.local, true)
+		if err != nil {
+			return err
+		}
+	}
+	return n.n.Close()
 }
 
 var _ Packet = (*nfqueuePacket)(nil)

ruleset/expr.go

@@ -46,7 +46,6 @@ type compiledExprRule struct {
 	Action      Action
 	ModInstance modifier.Instance
 	Program     *vm.Program
-	Analyzers   map[string]struct{}
 }
 
 var _ Ruleset = (*exprRuleset)(nil)
@@ -100,55 +99,45 @@ func CompileExprRules(rules []ExprRule, ans []analyzer.Analyzer, mods []modifier
 		if !ok {
 			return nil, fmt.Errorf("rule %q has invalid action %q", rule.Name, rule.Action)
 		}
-		visitor := &depVisitor{Analyzers: make(map[string]struct{})}
-		geoip := expr.Function(
-			"geoip",
-			func(params ...any) (any, error) {
-				return geoMatcher.MatchGeoIp(params[0].(string), params[1].(string)), nil
-			},
-			new(func(string, string) bool),
-		)
-		geosite := expr.Function(
-			"geosite",
-			func(params ...any) (any, error) {
-				return geoMatcher.MatchGeoSite(params[0].(string), params[1].(string)), nil
-			},
-			new(func(string, string) bool),
-		)
+		visitor := &idVisitor{Identifiers: make(map[string]bool)}
 		program, err := expr.Compile(rule.Expr,
 			func(c *conf.Config) {
 				c.Strict = false
 				c.Expect = reflect.Bool
 				c.Visitors = append(c.Visitors, visitor)
+				registerBuiltinFunctions(c.Functions, geoMatcher)
 			},
-			geoip,
-			geosite,
 		)
 		if err != nil {
 			return nil, fmt.Errorf("rule %q has invalid expression: %w", rule.Name, err)
 		}
-		for name := range visitor.Analyzers {
-			a, ok := fullAnMap[name]
-			if !ok && !isBuiltInAnalyzer(name) {
-				return nil, fmt.Errorf("rule %q uses unknown analyzer %q", rule.Name, name)
+		for name := range visitor.Identifiers {
+			if isBuiltInAnalyzer(name) {
+				continue
 			}
-			depAnMap[name] = a
-		}
-		if visitor.UseGeoSite {
-			if err := geoMatcher.LoadGeoSite(); err != nil {
-				return nil, fmt.Errorf("rule %q failed to load geosite: %w", rule.Name, err)
-			}
-		}
-		if visitor.UseGeoIp {
-			if err := geoMatcher.LoadGeoIP(); err != nil {
-				return nil, fmt.Errorf("rule %q failed to load geoip: %w", rule.Name, err)
+			// Check if it's one of the built-in functions, and if so,
+			// skip it as an analyzer & do initialization if necessary.
+			switch name {
+			case "geoip":
+				if err := geoMatcher.LoadGeoIP(); err != nil {
+					return nil, fmt.Errorf("rule %q failed to load geoip: %w", rule.Name, err)
+				}
+			case "geosite":
+				if err := geoMatcher.LoadGeoSite(); err != nil {
+					return nil, fmt.Errorf("rule %q failed to load geosite: %w", rule.Name, err)
+				}
+			default:
+				a, ok := fullAnMap[name]
+				if !ok {
+					return nil, fmt.Errorf("rule %q uses unknown analyzer %q", rule.Name, name)
+				}
+				depAnMap[name] = a
 			}
 		}
 		cr := compiledExprRule{
-			Name:      rule.Name,
-			Action:    action,
-			Program:   program,
-			Analyzers: visitor.Analyzers,
+			Name:    rule.Name,
+			Action:  action,
+			Program: program,
 		}
 		if action == ActionModify {
 			mod, ok := fullModMap[rule.Modifier.Name]
@@ -175,6 +164,23 @@ func CompileExprRules(rules []ExprRule, ans []analyzer.Analyzer, mods []modifier
 	}, nil
 }
 
+func registerBuiltinFunctions(funcMap map[string]*ast.Function, geoMatcher *geo.GeoMatcher) {
+	funcMap["geoip"] = &ast.Function{
+		Name: "geoip",
+		Func: func(params ...any) (any, error) {
+			return geoMatcher.MatchGeoIp(params[0].(string), params[1].(string)), nil
+		},
+		Types: []reflect.Type{reflect.TypeOf(geoMatcher.MatchGeoIp)},
+	}
+	funcMap["geosite"] = &ast.Function{
+		Name: "geosite",
+		Func: func(params ...any) (any, error) {
+			return geoMatcher.MatchGeoSite(params[0].(string), params[1].(string)), nil
+		},
+		Types: []reflect.Type{reflect.TypeOf(geoMatcher.MatchGeoSite)},
+	}
+}
+
 func streamInfoToExprEnv(info StreamInfo) map[string]interface{} {
 	m := map[string]interface{}{
 		"id":    info.ID,
@@ -241,22 +247,12 @@ func modifiersToMap(mods []modifier.Modifier) map[string]modifier.Modifier {
 	return modMap
 }
 
-type depVisitor struct {
-	Analyzers map[string]struct{}
-
-	UseGeoSite bool
-	UseGeoIp   bool
+type idVisitor struct {
+	Identifiers map[string]bool
 }
 
-func (v *depVisitor) Visit(node *ast.Node) {
+func (v *idVisitor) Visit(node *ast.Node) {
 	if idNode, ok := (*node).(*ast.IdentifierNode); ok {
-		switch idNode.Value {
-		case "geosite":
-			v.UseGeoSite = true
-		case "geoip":
-			v.UseGeoIp = true
-		default:
-			v.Analyzers[idNode.Value] = struct{}{}
-		}
+		v.Identifiers[idNode.Value] = true
 	}
 }