diff --git a/cmd/root.go b/cmd/root.go index aec1650..5821325 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -7,7 +7,6 @@ import ( "os/signal" "strconv" "strings" - "syscall" "github.com/apernet/OpenGFW/analyzer" "github.com/apernet/OpenGFW/analyzer/tcp" @@ -284,31 +283,13 @@ func runMain(cmd *cobra.Command, args []string) { 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") - } - } - }() + + rsLoader := ruleset.NewSignalRuleSetLoader(args[0], func(rs ruleset.Ruleset) error { + return en.UpdateRuleset(rs) + }, analyzers, modifiers, rsConfig) + if err := rsLoader.Start(); err != nil { + logger.Error("start ruleset loader failed", zap.Error(err)) + } logger.Info("engine started") logger.Info("engine exited", zap.Error(en.Run(ctx))) diff --git a/ruleset/loader.go b/ruleset/loader.go new file mode 100644 index 0000000..abb94ed --- /dev/null +++ b/ruleset/loader.go @@ -0,0 +1,72 @@ +package ruleset + +import ( + "os" + "os/signal" + "sync" + "syscall" + + "github.com/apernet/OpenGFW/analyzer" + "github.com/apernet/OpenGFW/modifier" + "go.uber.org/zap" +) + +type RuleSetHandler func(Ruleset) error + +type RuleSetLoader interface { + Start() error +} + +type SignalRuleSetLoader struct { + o sync.Once + filePath string + handler RuleSetHandler + reloadChan chan os.Signal + analyzers []analyzer.Analyzer + modifiers []modifier.Modifier + config *BuiltinConfig + logger *zap.Logger +} + +func NewSignalRuleSetLoader(path string, handler RuleSetHandler, ans []analyzer.Analyzer, mods []modifier.Modifier, cfg *BuiltinConfig) RuleSetLoader { + return &SignalRuleSetLoader{ + o: sync.Once{}, + filePath: path, + handler: handler, + reloadChan: make(chan os.Signal), + analyzers: ans, + modifiers: mods, + config: cfg, + } +} + +func (l *SignalRuleSetLoader) Start() error { + l.o.Do( + func() { + signal.Notify(l.reloadChan, syscall.SIGHUP) + go func() { + for { + <-l.reloadChan + l.logger.Info("reloading rules") + rawRs, err := ExprRulesFromYAML(l.filePath) + if err != nil { + l.logger.Error("failed to load rules, using old rules", zap.Error(err)) + continue + } + rs, err := CompileExprRules(rawRs, l.analyzers, l.modifiers, l.config) + if err != nil { + l.logger.Error("failed to compile rules, using old rules", zap.Error(err)) + continue + } + err = l.handler(rs) + if err != nil { + l.logger.Error("failed to update ruleset", zap.Error(err)) + } else { + l.logger.Info("rules reloaded") + } + } + }() + }, + ) + return nil +}