[minor] support regex patterns and suffix matching for protectRoutes

README.md

@@ -91,7 +91,8 @@ services:
 
 | **Parameter**           | **Type (Required)**     | **Default**              | **Description**                                                                                                                                                                                  |
 |-------------------------|-------------------------|--------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| `protectRoutes`         | `[]string` (required)   | `""`                     | Comma-separated list of route prefixes to protect. e.g., `"/"` protects the entire site (including file/js/css downloads, which you likely don't want). `"/browse"` protects its subtree.        |
+| `mode`                  | `string`                | `prefix`                 | Must be: `prefix`, `suffix`, `regex`. Matching does not include query parameters. `excludeRoutes` always uses `prefix` except when `mode: regex`. Only use `regex` when needed                   |
+| `protectRoutes`         | `[]string` (required)   | `""`                     | Comma-separated list of route prefixes/suffixes/regex patterns to protect.                                                                                                                       |
 | `excludeRoutes`         | `[]string`              | `""`                     | Comma-separated list of route prefixes to **never** protect. e.g., `protectRoutes: "/"` protects the entire site. `excludeRoutes: "/ajax"` would never challenge any route starting with `/ajax` |
 | `captchaProvider`       | `string` (required)     | `""`                     | The captcha type to use. Supported values: `turnstile`, `hcaptcha`, and `recaptcha`.                                                                                                             |
 | `siteKey`               | `string` (required)     | `""`                     | The captcha site key.                                                                                                                                                                            |
@@ -165,3 +166,50 @@ When you override the challenge template, the process probably looks like:
 - the original implementation of this logic was [a drupal module called turnstile_protect](https://www.drupal.org/project/turnstile_protect). This traefik plugin was made to make the challenge logic even more perfomant than that Drupal module, and also to provide this bot protection to non-Drupal websites
 - making general captcha structs to support multiple providers was based on the work in [crowdsec-bouncer-traefik-plugin](https://github.com/maxlerebourg/crowdsec-bouncer-traefik-plugin)
 - in memory cache thanks to https://github.com/patrickmn/go-cache
+
+## When to enable regex
+
+When possible, you want to keep regex disabled as seen in the example benchmark below.
+
+However, when needed it can be enabled with `mode: regex`
+
+```
+$ go mod init bench
+$ cat << EOF > bench_test.go
+package main
+
+import (
+	"regexp"
+	"strings"
+	"testing"
+)
+
+var (
+	testPath = "/api/v1/user/profile"
+	prefix   = "/api/v1"
+	regex    = regexp.MustCompile("^/api/v1")
+)
+
+func BenchmarkHasPrefix(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		_ = strings.HasPrefix(testPath, prefix)
+	}
+}
+
+func BenchmarkRegexMatch(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		_ = regex.MatchString(testPath)
+	}
+}
+EOF
+$ go test -bench=. -benchmem
+```
+
+```
+goos: darwin
+goarch: amd64
+cpu: Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz
+BenchmarkHasPrefix-12     	340856451	         3.415 ns/op	       0 B/op	       0 allocs/op
+BenchmarkRegexMatch-12    	27992568	        41.20 ns/op	       0 B/op	       0 allocs/op
+PASS
+```

main.go

@@ -10,6 +10,7 @@ import (
 	"net/url"
 	"os"
 	"path/filepath"
+	"regexp"
 	"slices"
 	"strings"
 	"text/template"
@@ -47,6 +48,9 @@ type Config struct {
 	EnableStatsPage       string   `json:"enableStatsPage"`
 	LogLevel              string   `json:"loglevel,omitempty"`
 	PersistentStateFile   string   `json:"persistentStateFile"`
+	Mode                  string   `json:"mode"`
+	protectRoutesRegex    []*regexp.Regexp
+	excludeRoutesRegex    []*regexp.Regexp
 }
 
 type CaptchaProtect struct {
@@ -93,6 +97,9 @@ func CreateConfig() *Config {
 		LogLevel:              "INFO",
 		IPDepth:               0,
 		CaptchaProvider:       "turnstile",
+		Mode:                  "prefix",
+		protectRoutesRegex:    []*regexp.Regexp{},
+		excludeRoutesRegex:    []*regexp.Regexp{},
 	}
 }
 
@@ -117,10 +124,29 @@ func NewCaptchaProtect(ctx context.Context, next http.Handler, config *Config, n
 	expiration := time.Duration(config.Window) * time.Second
 	log.Debug("Captcha config", "config", config)
 
-	if len(config.ProtectRoutes) == 0 {
+	if len(config.ProtectRoutes) == 0 && config.Mode != "suffix" {
 		return nil, fmt.Errorf("you must protect at least one route with the protectRoutes config value. / will cover your entire site")
 	}
 
+	if config.Mode == "regex" {
+		for _, r := range config.ProtectRoutes {
+			cr, err := regexp.Compile(r)
+			if err != nil {
+				return nil, fmt.Errorf("invalid regex in protectRoutes: %s", r)
+			}
+			config.protectRoutesRegex = append(config.protectRoutesRegex, cr)
+		}
+		for _, r := range config.ExcludeRoutes {
+			cr, err := regexp.Compile(r)
+			if err != nil {
+				return nil, fmt.Errorf("invalid regex in excludeRoutes: %s", r)
+			}
+			config.excludeRoutesRegex = append(config.excludeRoutesRegex, cr)
+		}
+	} else if config.Mode != "prefix" && config.Mode != "suffix" {
+		return nil, fmt.Errorf("unknown mode: %s. Supported values are prefix, suffix, and regex.", config.Mode)
+	}
+
 	if config.ChallengeURL == "/" {
 		return nil, fmt.Errorf("your challenge URL can not be the entire site. Default is `/challenge`. A blank value will have challenges presented on the visit that trips the rate limit")
 	}
@@ -398,10 +424,18 @@ func (bc *CaptchaProtect) shouldApply(req *http.Request, clientIP string) bool {
 		return false
 	}
 
-	return bc.RouteIsProtected(req.URL.Path)
+	if bc.config.Mode == "regex" {
+		return bc.RouteIsProtectedRegex(req.URL.Path)
+	}
+
+	if bc.config.Mode == "suffix" {
+		return bc.RouteIsProtectedSuffix(req.URL.Path)
+	}
+
+	return bc.RouteIsProtectedPrefix(req.URL.Path)
 }
 
-func (bc *CaptchaProtect) RouteIsProtected(path string) bool {
+func (bc *CaptchaProtect) RouteIsProtectedPrefix(path string) bool {
 protected:
 	for _, route := range bc.config.ProtectRoutes {
 		if !strings.HasPrefix(path, route) {
@@ -433,6 +467,73 @@ protected:
 	return false
 }
 
+func (bc *CaptchaProtect) RouteIsProtectedSuffix(path string) bool {
+protected:
+	for _, route := range bc.config.ProtectRoutes {
+		cleanPath := path
+		ext := filepath.Ext(path)
+		if ext != "" {
+			cleanPath = strings.TrimSuffix(path, ext)
+		}
+		if !strings.HasSuffix(cleanPath, route) {
+			continue
+		}
+
+		// we're on a protected route - make sure this route doesn't have an exclusion
+		for _, eRoute := range bc.config.ExcludeRoutes {
+			if strings.HasPrefix(cleanPath, eRoute) {
+				continue protected
+			}
+		}
+
+		// if this path isn't a file, go ahead and mark this path as protected
+		ext = strings.TrimPrefix(ext, ".")
+		if ext == "" {
+			return true
+		}
+
+		// if we have a file extension, see if we should protect this file extension type
+		for _, protectedExtensions := range bc.config.ProtectFileExtensions {
+			if strings.EqualFold(ext, protectedExtensions) {
+				return true
+			}
+		}
+	}
+
+	return false
+}
+
+func (bc *CaptchaProtect) RouteIsProtectedRegex(path string) bool {
+protected:
+	for _, routeRegex := range bc.config.protectRoutesRegex {
+		matched := routeRegex.MatchString(path)
+		if !matched {
+			continue
+		}
+
+		for _, excludeRegex := range bc.config.excludeRoutesRegex {
+			excluded := excludeRegex.MatchString(path)
+			if excluded {
+				continue protected
+			}
+		}
+
+		ext := filepath.Ext(path)
+		ext = strings.TrimPrefix(ext, ".")
+		if ext == "" {
+			return true
+		}
+
+		for _, protectedExtension := range bc.config.ProtectFileExtensions {
+			if strings.EqualFold(ext, protectedExtension) {
+				return true
+			}
+		}
+	}
+
+	return false
+}
+
 func IsIpExcluded(clientIP string, exemptIps []*net.IPNet) bool {
 	ip := net.ParseIP(clientIP)
 	for _, block := range exemptIps {

main_test.go

@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"regexp"
 	"strings"
 	"testing"
 )
@@ -389,20 +390,147 @@ func TestRouteIsProtected(t *testing.T) {
 		},
 	}
 
+	for _, tt := range tests {
+		for _, useRegex := range []bool{false, true} {
+			mode := "prefix"
+			if useRegex {
+				mode = "regex"
+			}
+
+			t.Run(tt.name+"_"+mode, func(t *testing.T) {
+				c := CreateConfig()
+				c.Mode = mode
+				c.ProtectFileExtensions = append(c.ProtectFileExtensions, tt.config.ProtectFileExtensions...)
+
+				if useRegex {
+					// Convert each route to ^... regex for "HasPrefix" behavior
+					for _, route := range tt.config.ProtectRoutes {
+						c.ProtectRoutes = append(c.ProtectRoutes, "^"+regexp.QuoteMeta(route))
+					}
+					for _, exclude := range tt.config.ExcludeRoutes {
+						c.ExcludeRoutes = append(c.ExcludeRoutes, "^"+regexp.QuoteMeta(exclude))
+					}
+				} else {
+					c.ProtectRoutes = append(c.ProtectRoutes, tt.config.ProtectRoutes...)
+					c.ExcludeRoutes = append(c.ExcludeRoutes, tt.config.ExcludeRoutes...)
+				}
+
+				bc, err := NewCaptchaProtect(context.Background(), nil, c, "captcha-protect")
+				if err != nil {
+					t.Errorf("unexpected error %v", err)
+				}
+
+				if useRegex {
+					result := bc.RouteIsProtectedRegex(tt.path)
+					if result != tt.expected {
+						t.Errorf("RouteIsProtected(%q) = %v; want %v (mode: %s)", tt.path, result, tt.expected, mode)
+					}
+				} else {
+					result := bc.RouteIsProtectedPrefix(tt.path)
+					if result != tt.expected {
+						t.Errorf("RouteIsProtected(%q) = %v; want %v (mode: %s)", tt.path, result, tt.expected, mode)
+					}
+				}
+			})
+		}
+	}
+
+}
+
+func TestRouteIsProtectedSuffix(t *testing.T) {
+	tests := []struct {
+		name     string
+		config   Config
+		path     string
+		expected bool
+	}{
+		{
+			name: "Protected route - suffix match",
+			config: Config{
+				ProtectRoutes:         []string{"baz"},
+				ProtectFileExtensions: []string{},
+				ExcludeRoutes:         []string{},
+			},
+			path:     "/baz",
+			expected: true,
+		},
+		{
+			name: "Unprotected route - no suffix match",
+			config: Config{
+				ProtectRoutes:         []string{"baz"},
+				ProtectFileExtensions: []string{},
+				ExcludeRoutes:         []string{},
+			},
+			path:     "/foo/bar",
+			expected: false,
+		},
+		{
+			name: "Protected route with file extension",
+			config: Config{
+				ProtectRoutes:         []string{"style"},
+				ProtectFileExtensions: []string{"json"},
+				ExcludeRoutes:         []string{},
+			},
+			path:     "/foo/bar/style.json",
+			expected: true,
+		},
+		{
+			name: "Protected by extension only",
+			config: Config{
+				ProtectRoutes:         []string{"index"},
+				ProtectFileExtensions: []string{"html"},
+				ExcludeRoutes:         []string{},
+			},
+			path:     "/whatever/index.html",
+			expected: true,
+		},
+		{
+			name: "Suffix protected route",
+			config: Config{
+				ProtectRoutes:         []string{"route"},
+				ProtectFileExtensions: []string{},
+				ExcludeRoutes:         []string{},
+			},
+			path:     "/bar/any/route",
+			expected: true,
+		},
+		{
+			name: "Excluded route - suffix match",
+			config: Config{
+				ProtectRoutes:         []string{"ajax"},
+				ProtectFileExtensions: []string{},
+				ExcludeRoutes:         []string{"/api"},
+			},
+			path:     "/api/ajax",
+			expected: false,
+		},
+		{
+			name: "Protected route - not excluded",
+			config: Config{
+				ProtectRoutes:         []string{"ajax"},
+				ProtectFileExtensions: []string{},
+				ExcludeRoutes:         []string{"notajax"},
+			},
+			path:     "/real/ajax",
+			expected: true,
+		},
+	}
+
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			c := CreateConfig()
 			c.ProtectRoutes = append(c.ProtectRoutes, tt.config.ProtectRoutes...)
 			c.ExcludeRoutes = append(c.ExcludeRoutes, tt.config.ExcludeRoutes...)
+			c.Mode = "suffix"
 			c.ProtectFileExtensions = append(c.ProtectFileExtensions, tt.config.ProtectFileExtensions...)
 			bc, err := NewCaptchaProtect(context.Background(), nil, c, "captcha-protect")
 			if err != nil {
 				t.Errorf("unexpected error %v", err)
 			}
 
-			result := bc.RouteIsProtected(tt.path)
+			result := bc.RouteIsProtectedSuffix(tt.path)
 			if result != tt.expected {
-				t.Errorf("RouteIsProtected(%q) = %v; want %v", tt.path, result, tt.expected)
+				t.Errorf("RouteIsProtectedSuffix(%q) = %v; want %v", tt.path, result, tt.expected)
 			}
 		})
 	}