From 6e964e64491501ddfc8efe84ad869fc7f55910ee Mon Sep 17 00:00:00 2001
From: Xe Iaso <me@xeiaso.net>
Date: Thu, 15 May 2025 08:29:37 -0400
Subject: [PATCH] feat(config): add WEIGH action allowing admins to adjust
 request weight

Signed-off-by: Xe Iaso <me@xeiaso.net>
---
 lib/policy/config/config.go                   | 18 ++++++++++++------
 lib/policy/config/config_test.go              | 19 +++++++++++++++++++
 lib/policy/config/expressionorlist.go         |  4 ++--
 .../config/testdata/good/simple-weight.yaml   |  6 ++++++
 .../testdata/good/weight-no-weight.yaml       |  4 ++++
 lib/policy/config/weight.go                   |  5 +++++
 6 files changed, 48 insertions(+), 8 deletions(-)
 create mode 100644 lib/policy/config/testdata/good/simple-weight.yaml
 create mode 100644 lib/policy/config/testdata/good/weight-no-weight.yaml
 create mode 100644 lib/policy/config/weight.go

diff --git a/lib/policy/config/config.go b/lib/policy/config/config.go
index 4b6f6435..86b59acc 100644
--- a/lib/policy/config/config.go
+++ b/lib/policy/config/config.go
@@ -39,6 +39,7 @@ const (
 	RuleAllow     Rule = "ALLOW"
 	RuleDeny      Rule = "DENY"
 	RuleChallenge Rule = "CHALLENGE"
+	RuleWeigh     Rule = "WEIGH"
 	RuleBenchmark Rule = "DEBUG_BENCHMARK"
 )
 
@@ -51,14 +52,15 @@ const (
 )
 
 type BotConfig struct {
-	UserAgentRegex *string           `json:"user_agent_regex"`
-	PathRegex      *string           `json:"path_regex"`
-	HeadersRegex   map[string]string `json:"headers_regex"`
-	Expression     *ExpressionOrList `json:"expression"`
+	UserAgentRegex *string           `json:"user_agent_regex,omitempty"`
+	PathRegex      *string           `json:"path_regex,omitempty"`
+	HeadersRegex   map[string]string `json:"headers_regex,omitempty"`
+	Expression     *ExpressionOrList `json:"expression,omitempty"`
 	Challenge      *ChallengeRules   `json:"challenge,omitempty"`
+	Weight         *Weight           `json:"weight,omitempty"`
 	Name           string            `json:"name"`
 	Action         Rule              `json:"action"`
-	RemoteAddr     []string          `json:"remote_addresses"`
+	RemoteAddr     []string          `json:"remote_addresses,omitempty"`
 }
 
 func (b BotConfig) Zero() bool {
@@ -150,7 +152,7 @@ func (b BotConfig) Valid() error {
 	}
 
 	switch b.Action {
-	case RuleAllow, RuleBenchmark, RuleChallenge, RuleDeny:
+	case RuleAllow, RuleBenchmark, RuleChallenge, RuleDeny, RuleWeigh:
 		// okay
 	default:
 		errs = append(errs, fmt.Errorf("%w: %q", ErrUnknownAction, b.Action))
@@ -162,6 +164,10 @@ func (b BotConfig) Valid() error {
 		}
 	}
 
+	if b.Action == RuleWeigh && b.Weight == nil {
+		b.Weight = &Weight{Adjust: 5}
+	}
+
 	if len(errs) != 0 {
 		return fmt.Errorf("config: bot entry for %q is not valid:\n%w", b.Name, errors.Join(errs...))
 	}
diff --git a/lib/policy/config/config_test.go b/lib/policy/config/config_test.go
index afc1ab88..590f7fce 100644
--- a/lib/policy/config/config_test.go
+++ b/lib/policy/config/config_test.go
@@ -182,6 +182,25 @@ func TestBotValid(t *testing.T) {
 			},
 			err: nil,
 		},
+		{
+			name: "weight rule without weight",
+			bot: BotConfig{
+				Name:           "weight-adjust-if-mozilla",
+				Action:         RuleWeigh,
+				UserAgentRegex: p("Mozilla"),
+			},
+		},
+		{
+			name: "weight rule with weight adjust",
+			bot: BotConfig{
+				Name:           "weight-adjust-if-mozilla",
+				Action:         RuleWeigh,
+				UserAgentRegex: p("Mozilla"),
+				Weight: &Weight{
+					Adjust: 5,
+				},
+			},
+		},
 	}
 
 	for _, cs := range tests {
diff --git a/lib/policy/config/expressionorlist.go b/lib/policy/config/expressionorlist.go
index 7b07a359..6f4a8904 100644
--- a/lib/policy/config/expressionorlist.go
+++ b/lib/policy/config/expressionorlist.go
@@ -14,8 +14,8 @@ var (
 
 type ExpressionOrList struct {
 	Expression string   `json:"-"`
-	All        []string `json:"all"`
-	Any        []string `json:"any"`
+	All        []string `json:"all,omitempty"`
+	Any        []string `json:"any,omitempty"`
 }
 
 func (eol ExpressionOrList) Equal(rhs *ExpressionOrList) bool {
diff --git a/lib/policy/config/testdata/good/simple-weight.yaml b/lib/policy/config/testdata/good/simple-weight.yaml
new file mode 100644
index 00000000..79702d4f
--- /dev/null
+++ b/lib/policy/config/testdata/good/simple-weight.yaml
@@ -0,0 +1,6 @@
+bots:
+- name: simple-weight-adjust
+  action: WEIGH
+  user_agent_regex: Mozilla
+  weight:
+    adjust: 5
\ No newline at end of file
diff --git a/lib/policy/config/testdata/good/weight-no-weight.yaml b/lib/policy/config/testdata/good/weight-no-weight.yaml
new file mode 100644
index 00000000..dfdb201d
--- /dev/null
+++ b/lib/policy/config/testdata/good/weight-no-weight.yaml
@@ -0,0 +1,4 @@
+bots:
+- name: weight
+  action: WEIGH
+  user_agent_regex: Mozilla
\ No newline at end of file
diff --git a/lib/policy/config/weight.go b/lib/policy/config/weight.go
new file mode 100644
index 00000000..f408111b
--- /dev/null
+++ b/lib/policy/config/weight.go
@@ -0,0 +1,5 @@
+package config
+
+type Weight struct {
+	Adjust int `json:"adjust"`
+}