support probabilistic packet drop for tcp

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

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)
 

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

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:

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

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 {

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"