diff --git a/cmd/root.go b/cmd/root.go index ed16c34..0041a4e 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -348,6 +348,15 @@ func (l *engineLogger) TCPStreamAction(info ruleset.StreamInfo, action ruleset.A zap.Bool("noMatch", noMatch)) } +func (l *engineLogger) TCPStreamActionType(info ruleset.StreamInfo, action ruleset.ActionType, noMatch bool) { + logger.Info("TCP stream action", + zap.Int64("id", info.ID), + zap.String("src", info.SrcString()), + zap.String("dst", info.DstString()), + zap.String("action", action.String()), + zap.Bool("noMatch", noMatch)) +} + func (l *engineLogger) UDPStreamNew(workerID int, info ruleset.StreamInfo) { logger.Debug("new UDP stream", zap.Int("workerID", workerID), @@ -374,6 +383,15 @@ func (l *engineLogger) UDPStreamAction(info ruleset.StreamInfo, action ruleset.A zap.Bool("noMatch", noMatch)) } +func (l *engineLogger) UDPStreamActionType(info ruleset.StreamInfo, action ruleset.ActionType, noMatch bool) { + logger.Info("UDP stream action", + zap.Int64("id", info.ID), + zap.String("src", info.SrcString()), + zap.String("dst", info.DstString()), + zap.String("action", action.String()), + zap.Bool("noMatch", noMatch)) +} + func (l *engineLogger) ModifyError(info ruleset.StreamInfo, err error) { logger.Error("modify error", zap.Int64("id", info.ID), diff --git a/engine/interface.go b/engine/interface.go index b5a4f57..4dad38c 100644 --- a/engine/interface.go +++ b/engine/interface.go @@ -36,10 +36,12 @@ type Logger interface { TCPStreamNew(workerID int, info ruleset.StreamInfo) TCPStreamPropUpdate(info ruleset.StreamInfo, close bool) TCPStreamAction(info ruleset.StreamInfo, action ruleset.Action, noMatch bool) + TCPStreamActionType(info ruleset.StreamInfo, actionType ruleset.ActionType, noMatch bool) UDPStreamNew(workerID int, info ruleset.StreamInfo) UDPStreamPropUpdate(info ruleset.StreamInfo, close bool) UDPStreamAction(info ruleset.StreamInfo, action ruleset.Action, noMatch bool) + UDPStreamActionType(info ruleset.StreamInfo, actionType ruleset.ActionType, noMatch bool) ModifyError(info ruleset.StreamInfo, err error) diff --git a/engine/tcp.go b/engine/tcp.go index 4bf776c..50a2b76 100644 --- a/engine/tcp.go +++ b/engine/tcp.go @@ -22,6 +22,7 @@ type tcpVerdict io.Verdict const ( tcpVerdictAccept = tcpVerdict(io.VerdictAccept) tcpVerdictAcceptStream = tcpVerdict(io.VerdictAcceptStream) + tcpVerdictDrop = tcpVerdict(io.VerdictDrop) tcpVerdictDropStream = tcpVerdict(io.VerdictDropStream) ) @@ -103,6 +104,11 @@ type tcpStream struct { activeEntries []*tcpStreamEntry doneEntries []*tcpStreamEntry lastVerdict tcpVerdict + + // 0: probabilistic drop disabled + // 1-99: drop each packet with this probability + // 100: drop stream + dropRate uint8 } type tcpStreamEntry struct { @@ -120,6 +126,15 @@ func (s *tcpStream) Accept(tcp *layers.TCP, ci gopacket.CaptureInfo, dir reassem return true } else { ctx := ac.(*tcpContext) + // Probability drop mode + if s.dropRate > 0 { + if shouldDrop(s.dropRate) { + ctx.Verdict = tcpVerdictDrop + } else { + ctx.Verdict = tcpVerdictAccept + } + return false + } ctx.Verdict = s.lastVerdict return false } @@ -150,7 +165,23 @@ func (s *tcpStream) ReassembledSG(sg reassembly.ScatterGather, ac reassembly.Ass // Match properties against ruleset result := s.ruleset.Match(s.info) action := result.Action - if action != ruleset.ActionMaybe && action != ruleset.ActionModify { + if action.Type == ruleset.ActionDrop { + rate := *action.Params.DropRate + if rate == 100 { + s.lastVerdict = tcpVerdictDropStream + ctx.Verdict = tcpVerdictDropStream + } else if rate == 0 { + s.lastVerdict = tcpVerdictAcceptStream + ctx.Verdict = tcpVerdictAcceptStream + } else { + // Enter drop mode + s.dropRate = uint8(rate) + ctx.Verdict = tcpVerdictAccept + } + s.logger.TCPStreamAction(s.info, action, false) + // No more analyze + s.closeActiveEntries() + } else if action.Type != ruleset.ActionMaybe && action.Type != ruleset.ActionModify { verdict := actionToTCPVerdict(action) s.lastVerdict = verdict ctx.Verdict = verdict @@ -159,11 +190,11 @@ func (s *tcpStream) ReassembledSG(sg reassembly.ScatterGather, ac reassembly.Ass s.closeActiveEntries() } } - if len(s.activeEntries) == 0 && ctx.Verdict == tcpVerdictAccept { + if len(s.activeEntries) == 0 && ctx.Verdict == tcpVerdictAccept && s.dropRate == 0 { // All entries are done but no verdict issued, accept stream s.lastVerdict = tcpVerdictAcceptStream ctx.Verdict = tcpVerdictAcceptStream - s.logger.TCPStreamAction(s.info, ruleset.ActionAllow, true) + s.logger.TCPStreamActionType(s.info, ruleset.ActionAllow, true) } } @@ -214,11 +245,18 @@ func analyzersToTCPAnalyzers(ans []analyzer.Analyzer) []analyzer.TCPAnalyzer { } func actionToTCPVerdict(a ruleset.Action) tcpVerdict { - switch a { + switch a.Type { case ruleset.ActionMaybe, ruleset.ActionAllow, ruleset.ActionModify: return tcpVerdictAcceptStream - case ruleset.ActionBlock, ruleset.ActionDrop: + case ruleset.ActionBlock: return tcpVerdictDropStream + case ruleset.ActionDrop: + // Should never happen + rate := *a.Params.DropRate + if rate == 100 { + return tcpVerdictDropStream + } + return tcpVerdictAccept default: // Should never happen return tcpVerdictAcceptStream diff --git a/engine/udp.go b/engine/udp.go index fe1f07c..ce8d9ce 100644 --- a/engine/udp.go +++ b/engine/udp.go @@ -203,24 +203,24 @@ func (s *udpStream) Feed(udp *layers.UDP, rev bool, uc *udpContext) { // Match properties against ruleset result := s.ruleset.Match(s.info) action := result.Action - if action == ruleset.ActionModify { + if action.Type == ruleset.ActionModify { // Call the modifier instance udpMI, ok := result.ModInstance.(modifier.UDPModifierInstance) if !ok { // Not for UDP, fallback to maybe s.logger.ModifyError(s.info, errInvalidModifier) - action = ruleset.ActionMaybe + action.Type = ruleset.ActionMaybe } else { var err error uc.Packet, err = udpMI.Process(udp.Payload) if err != nil { // Modifier error, fallback to maybe s.logger.ModifyError(s.info, err) - action = ruleset.ActionMaybe + action.Type = ruleset.ActionMaybe } } } - if action != ruleset.ActionMaybe { + if action.Type != ruleset.ActionMaybe { verdict, final := actionToUDPVerdict(action) s.lastVerdict = verdict uc.Verdict = verdict @@ -234,7 +234,7 @@ func (s *udpStream) Feed(udp *layers.UDP, rev bool, uc *udpContext) { // All entries are done but no verdict issued, accept stream s.lastVerdict = udpVerdictAcceptStream uc.Verdict = udpVerdictAcceptStream - s.logger.UDPStreamAction(s.info, ruleset.ActionAllow, true) + s.logger.UDPStreamActionType(s.info, ruleset.ActionAllow, true) } } @@ -281,7 +281,7 @@ func analyzersToUDPAnalyzers(ans []analyzer.Analyzer) []analyzer.UDPAnalyzer { } func actionToUDPVerdict(a ruleset.Action) (v udpVerdict, final bool) { - switch a { + switch a.Type { case ruleset.ActionMaybe: return udpVerdictAccept, false case ruleset.ActionAllow: diff --git a/engine/utils.go b/engine/utils.go index 3b08c83..b2e37d2 100644 --- a/engine/utils.go +++ b/engine/utils.go @@ -1,6 +1,10 @@ package engine -import "github.com/apernet/OpenGFW/analyzer" +import ( + "math/rand/v2" + + "github.com/apernet/OpenGFW/analyzer" +) var _ analyzer.Logger = (*analyzerLogger)(nil) @@ -48,3 +52,7 @@ func processPropUpdate(cpm analyzer.CombinedPropMap, name string, update *analyz return false } } + +func shouldDrop(rate uint8) bool { + return rand.IntN(100) < int(rate) +} diff --git a/ruleset/expr.go b/ruleset/expr.go index 29d57c2..255b048 100644 --- a/ruleset/expr.go +++ b/ruleset/expr.go @@ -25,11 +25,16 @@ import ( // ExprRule is the external representation of an expression rule. type ExprRule struct { - Name string `yaml:"name"` - Action string `yaml:"action"` - Log bool `yaml:"log"` - Modifier ModifierEntry `yaml:"modifier"` - Expr string `yaml:"expr"` + Name string `yaml:"name"` + Action string `yaml:"action"` + ActionParams *ActionParams `yaml:"actionParams,omitempty"` + Log bool `yaml:"log"` + Modifier ModifierEntry `yaml:"modifier"` + Expr string `yaml:"expr"` +} + +type ActionParams struct { + DropRate *int `yaml:"dropRate,omitempty"` } type ModifierEntry struct { @@ -91,7 +96,7 @@ func (r *exprRuleset) Match(info StreamInfo) MatchResult { } // No match return MatchResult{ - Action: ActionMaybe, + Action: Action{Type: ActionMaybe}, } } @@ -111,11 +116,15 @@ func CompileExprRules(rules []ExprRule, ans []analyzer.Analyzer, mods []modifier } var action *Action if rule.Action != "" { - a, ok := actionStringToAction(rule.Action) + a, ok := actionStringToActionType(rule.Action) if !ok { return nil, fmt.Errorf("rule %q has invalid action %q", rule.Name, rule.Action) } - action = &a + params, err := normalizeActionParams(&rule, a) + if err != nil { + return nil, err + } + action = &Action{Type: a, Params: params} } visitor := &idVisitor{Variables: make(map[string]bool), Identifiers: make(map[string]bool)} patcher := &idPatcher{FuncMap: funcMap} @@ -162,7 +171,7 @@ func CompileExprRules(rules []ExprRule, ans []analyzer.Analyzer, mods []modifier Log: rule.Log, Program: program, } - if action != nil && *action == ActionModify { + if action != nil && action.Type == ActionModify { mod, ok := fullModMap[rule.Modifier.Name] if !ok { return nil, fmt.Errorf("rule %q uses unknown modifier %q", rule.Name, rule.Modifier.Name) @@ -218,7 +227,7 @@ func isBuiltInAnalyzer(name string) bool { } } -func actionStringToAction(action string) (Action, bool) { +func actionStringToActionType(action string) (ActionType, bool) { switch strings.ToLower(action) { case "allow": return ActionAllow, true @@ -253,6 +262,36 @@ func modifiersToMap(mods []modifier.Modifier) map[string]modifier.Modifier { return modMap } +func normalizeActionParams(rule *ExprRule, actionType ActionType) (*ActionParams, error) { + params := &ActionParams{} + if rule.ActionParams != nil { + *params = *rule.ActionParams + } + if actionType != ActionDrop { + if params.DropRate != nil { + return nil, fmt.Errorf( + "rule %q: dropRate is only valid for drop action", + rule.Name, + ) + } + return params, nil + } + // dropRate defaults to 100 for backward compatibility. + if params.DropRate == nil { + rate := 100 + params.DropRate = &rate + } + rate := *params.DropRate + if rate < 0 || rate > 100 { + return nil, fmt.Errorf( + "rule %q: dropRate must be between 0 and 100, got %d", + rule.Name, + rate, + ) + } + return params, nil +} + // idVisitor is a visitor that collects all identifiers in an expression. // This is for determining which analyzers are used by the expression. type idVisitor struct { diff --git a/ruleset/interface.go b/ruleset/interface.go index 535c2a4..eac786a 100644 --- a/ruleset/interface.go +++ b/ruleset/interface.go @@ -9,12 +9,17 @@ import ( "github.com/apernet/OpenGFW/modifier" ) -type Action int +type Action struct { + Type ActionType + Params *ActionParams +} + +type ActionType int const ( // ActionMaybe indicates that the ruleset hasn't seen anything worth blocking based on // current information, but that may change if volatile fields change in the future. - ActionMaybe Action = iota + ActionMaybe ActionType = iota // ActionAllow indicates that the stream should be allowed regardless of future changes. ActionAllow // ActionBlock indicates that the stream should be blocked. @@ -29,7 +34,13 @@ const ( ActionModify ) -func (a Action) String() string { +func (a *Action) String() string { + ret := a.Type.String() + // todo: concat parameters string + return ret +} + +func (a ActionType) String() string { switch a { case ActionMaybe: return "maybe"