Add headers bot rule (#300)

* Closes #291: add headers support to bot policy rules

* Fix config validator

7 files changed, 114 insertions, 11 deletions

diff --git a/lib/anubis.go b/lib/anubis.go
index 353fe63..ba143f9 100644
--- a/lib/anubis.go
+++ b/lib/anubis.go
@@ -548,6 +548,12 @@ func (s *Server) check(r *http.Request) (CheckResult, *policy.Bot, error) {
 				return cr("bot/"+b.Name, b.Action), &b, nil
 			}
 		}
+
+		if len(b.Headers) > 0 {
+			if s.checkHeaders(b, r.Header) {
+				return cr("bot/"+b.Name, b.Action), &b, nil
+			}
+		}
 	}
 
 	return cr("default/allow", config.RuleAllow), &policy.Bot{
@@ -572,6 +578,27 @@ func (s *Server) checkRemoteAddress(b policy.Bot, addr net.IP) bool {
 	return ok
 }
 
+func (s *Server) checkHeaders(b policy.Bot, header http.Header) bool {
+	if len(b.Headers) == 0 {
+		return true
+	}
+
+	for name, expr := range b.Headers {
+		values := header.Values(name)
+		if values == nil {
+			return false
+		}
+
+		for _, value := range values {
+			if !expr.MatchString(value) {
+				return false
+			}
+		}
+	}
+
+	return true
+}
+
 func (s *Server) CleanupDecayMap() {
 	s.DNSBLCache.Cleanup()
 	s.OGTags.Cleanup()
diff --git a/lib/policy/bot.go b/lib/policy/bot.go
index d9ca135..e656d9a 100644
--- a/lib/policy/bot.go
+++ b/lib/policy/bot.go
@@ -3,6 +3,7 @@ package policy
 import (
 	"fmt"
 	"regexp"
+	"strings"
 
 	"github.com/TecharoHQ/anubis/internal"
 	"github.com/TecharoHQ/anubis/lib/policy/config"
@@ -13,6 +14,7 @@ type Bot struct {
 	Name      string
 	UserAgent *regexp.Regexp
 	Path      *regexp.Regexp
+	Headers   map[string]*regexp.Regexp
 	Action    config.Rule `json:"action"`
 	Challenge *config.ChallengeRules
 	Ranger    cidranger.Ranger
@@ -27,6 +29,18 @@ func (b Bot) Hash() (string, error) {
 	if b.UserAgent != nil {
 		userAgentRex = b.UserAgent.String()
 	}
+	var headersRex string
+	if len(b.Headers) > 0 {
+		var sb strings.Builder
+		sb.Grow(len(b.Headers) * 64)
 
-	return internal.SHA256sum(fmt.Sprintf("%s::%s::%s", b.Name, pathRex, userAgentRex)), nil
+		for name, expr := range b.Headers {
+			sb.WriteString(name)
+			sb.WriteString(expr.String())
+		}
+
+		headersRex = sb.String()
+	}
+
+	return internal.SHA256sum(fmt.Sprintf("%s::%s::%s::%s", b.Name, pathRex, userAgentRex, headersRex)), nil
 }
diff --git a/lib/policy/config/config.go b/lib/policy/config/config.go
index e8f5161..b3d5cac 100644
--- a/lib/policy/config/config.go
+++ b/lib/policy/config/config.go
@@ -10,11 +10,12 @@ import (
 var (
 	ErrNoBotRulesDefined                 = errors.New("config: must define at least one (1) bot rule")
 	ErrBotMustHaveName                   = errors.New("config.Bot: must set name")
-	ErrBotMustHaveUserAgentOrPath        = errors.New("config.Bot: must set either user_agent_regex, path_regex, or remote_addresses")
+	ErrBotMustHaveUserAgentOrPath        = errors.New("config.Bot: must set either user_agent_regex, path_regex, headers_regex, or remote_addresses")
 	ErrBotMustHaveUserAgentOrPathNotBoth = errors.New("config.Bot: must set either user_agent_regex, path_regex, and not both")
 	ErrUnknownAction                     = errors.New("config.Bot: unknown action")
 	ErrInvalidUserAgentRegex             = errors.New("config.Bot: invalid user agent regex")
 	ErrInvalidPathRegex                  = errors.New("config.Bot: invalid path regex")
+	ErrInvalidHeadersRegex               = errors.New("config.Bot: invalid headers regex")
 	ErrInvalidCIDR                       = errors.New("config.Bot: invalid CIDR")
 )
 
@@ -37,12 +38,13 @@ const (
 )
 
 type BotConfig struct {
-	Name           string          `json:"name"`
-	UserAgentRegex *string         `json:"user_agent_regex"`
-	PathRegex      *string         `json:"path_regex"`
-	Action         Rule            `json:"action"`
-	RemoteAddr     []string        `json:"remote_addresses"`
-	Challenge      *ChallengeRules `json:"challenge,omitempty"`
+	Name           string            `json:"name"`
+	UserAgentRegex *string           `json:"user_agent_regex"`
+	PathRegex      *string           `json:"path_regex"`
+	HeadersRegex   map[string]string `json:"headers_regex"`
+	Action         Rule              `json:"action"`
+	RemoteAddr     []string          `json:"remote_addresses"`
+	Challenge      *ChallengeRules   `json:"challenge,omitempty"`
 }
 
 func (b BotConfig) Valid() error {
@@ -52,7 +54,7 @@ func (b BotConfig) Valid() error {
 		errs = append(errs, ErrBotMustHaveName)
 	}
 
-	if b.UserAgentRegex == nil && b.PathRegex == nil && len(b.RemoteAddr) == 0 {
+	if b.UserAgentRegex == nil && b.PathRegex == nil && len(b.RemoteAddr) == 0 && len(b.HeadersRegex) == 0 {
 		errs = append(errs, ErrBotMustHaveUserAgentOrPath)
 	}
 
@@ -72,6 +74,18 @@ func (b BotConfig) Valid() error {
 		}
 	}
 
+	if len(b.HeadersRegex) > 0 {
+		for name, expr := range b.HeadersRegex {
+			if name == "" {
+				continue
+			}
+
+			if _, err := regexp.Compile(expr); err != nil {
+				errs = append(errs, ErrInvalidHeadersRegex, err)
+			}
+		}
+	}
+
 	if len(b.RemoteAddr) > 0 {
 		for _, cidr := range b.RemoteAddr {
 			if _, _, err := net.ParseCIDR(cidr); err != nil {
diff --git a/lib/policy/config/config_test.go b/lib/policy/config/config_test.go
index a169087..0fabbb7 100644
--- a/lib/policy/config/config_test.go
+++ b/lib/policy/config/config_test.go
@@ -88,6 +88,18 @@ func TestBotValid(t *testing.T) {
 			err: ErrInvalidPathRegex,
 		},
 		{
+			name: "invalid headers regex",
+			bot: BotConfig{
+				Name:   "mozilla-ua",
+				Action: RuleChallenge,
+				HeadersRegex: map[string]string{
+					"Content-Type": "a(b",
+				},
+				PathRegex: p("a(b"),
+			},
+			err: ErrInvalidHeadersRegex,
+		},
+		{
 			name: "challenge difficulty too low",
 			bot: BotConfig{
 				Name:      "mozilla-ua",
diff --git a/lib/policy/config/testdata/bad/badregexes.json b/lib/policy/config/testdata/bad/badregexes.json
index e85b85b..db371b0 100644
--- a/lib/policy/config/testdata/bad/badregexes.json
+++ b/lib/policy/config/testdata/bad/badregexes.json
@@ -9,6 +9,13 @@
       "name": "user-agent-bad",
       "user_agent_regex": "a(b",
       "action": "DENY"
+    },
+    {
+      "name": "headers-bad",
+      "headers": {
+        "Accept-Encoding": "a(b"
+      },
+      "action": "DENY"
     }
   ]
 }
\ No newline at end of file
diff --git a/lib/policy/config/testdata/good/block_cf_workers.json b/lib/policy/config/testdata/good/block_cf_workers.json
new file mode 100644
index 0000000..b84f1e0
--- /dev/null
+++ b/lib/policy/config/testdata/good/block_cf_workers.json
@@ -0,0 +1,12 @@
+{
+  "bots": [
+    {
+      "name": "Cloudflare Workers",
+      "headers_regex": {
+        "CF-Worker": ".*"
+      },
+      "action": "DENY"
+    }
+  ],
+  "dnsbl": false
+}
\ No newline at end of file
diff --git a/lib/policy/policy.go b/lib/policy/policy.go
index 3f80fa3..4451b08 100644
--- a/lib/policy/policy.go
+++ b/lib/policy/policy.go
@@ -58,8 +58,9 @@ func ParseConfig(fin io.Reader, fname string, defaultDifficulty int) (*ParsedCon
 		}
 
 		parsedBot := Bot{
-			Name:   b.Name,
-			Action: b.Action,
+			Name:    b.Name,
+			Action:  b.Action,
+			Headers: map[string]*regexp.Regexp{},
 		}
 
 		if len(b.RemoteAddr) > 0 {
@@ -95,6 +96,22 @@ func ParseConfig(fin io.Reader, fname string, defaultDifficulty int) (*ParsedCon
 			}
 		}
 
+		if len(b.HeadersRegex) > 0 {
+			for name, expr := range b.HeadersRegex {
+				if name == "" {
+					continue
+				}
+
+				header, err := regexp.Compile(expr)
+				if err != nil {
+					validationErrs = append(validationErrs, fmt.Errorf("while compiling header regexp: %w", err))
+					continue
+				} else {
+					parsedBot.Headers[name] = header
+				}
+			}
+		}
+
 		if b.Challenge == nil {
 			parsedBot.Challenge = &config.ChallengeRules{
 				Difficulty: defaultDifficulty,