feat: rule set loader

Add rule set loader, move original config loader routine into
ruleset/loader.go
This commit is contained in:
Aionw 2024-03-02 15:16:17 +08:00
parent 1dce82745d
commit 65cb64168d
2 changed files with 79 additions and 26 deletions

View File

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

72
ruleset/loader.go Normal file
View File

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