[minor] persist state across container restarts

README.md

@@ -32,6 +32,7 @@ services:
             traefik.http.middlewares.captcha-protect.plugin.captcha-protect.siteKey: ${TURNSTILE_SITE_KEY}
             traefik.http.middlewares.captcha-protect.plugin.captcha-protect.secretKey: ${TURNSTILE_SECRET_KEY}
             traefik.http.middlewares.captcha-protect.plugin.captcha-protect.goodBots: apple.com,archive.org,duckduckgo.com,facebook.com,google.com,googlebot.com,googleusercontent.com,instagram.com,kagibot.org,linkedin.com,msn.com,openalex.org,twitter.com,x.com
+            traefik.http.middlewares.captcha-protect.plugin.captcha-protect.persistentStateFile: /tmp/state.json
         networks:
             default:
                 aliases:
@@ -50,6 +51,7 @@ services:
             --experimental.plugins.captcha-protect.version=v1.0.0
         volumes:
             - /var/run/docker.sock:/var/run/docker.sock:z
+            - /CHANGEME/TO/A/HOST/PATH/FOR/STATE/FILE:/tmp/state.json:rw
         ports:
             - "80:80"
         networks:
@@ -80,8 +82,9 @@ services:
 | exemptIps           | []string              | privateIPs              | IP address(es) in CIDR format that should never be challenged. Private IP ranges are always included                                                               |
 | challengeURL        | string                | "/challenge"            | The URL on the site to send challenges to. Will override any URL at that route                                                                                     |
 | challengeTmpl       | string                | "./challenge.tmpl.html" | HTML go template file to serve the captcha challenge.                                                                                                              |
-| enableStatsPage     | string                | "false"                 | Allow `exemptIps` to access `/captcha-protect/stats` to see the status of the rate limiter                                                                           |
+| enableStatsPage     | string                | "false"                 | Allow `exemptIps` to access `/captcha-protect/stats` to see the status of the rate limiter                                                                         |
 | logLevel            | string                | "INFO"                  | This middleware's log level. Possible values: ERROR, WARNING, INFO, or DEBUG                                                                                       |
+| persistentStateFile | string                | ""                      | When traefik restarts, the rate limit is reset. You can save the state to a file and have it reload on restart. In docker, requires mounting a file from the host. |
 
 
 ### Good Bots

ci/.gitignore

@@ -0,0 +1,2 @@
+tmp/*
+!tmp/.keep

ci/docker-compose.yml

@@ -21,6 +21,7 @@ services:
             traefik.http.middlewares.captcha-protect.plugin.captcha-protect.logLevel: "DEBUG"
             traefik.http.middlewares.captcha-protect.plugin.captcha-protect.goodBots: ""
             traefik.http.middlewares.captcha-protect.plugin.captcha-protect.protectRoutes: "/"
+            traefik.http.middlewares.captcha-protect.plugin.captcha-protect.persistentStateFile: "/tmp/state.json"
         healthcheck:
             test: curl -fs http://localhost/healthz | grep -q OK || exit 1
         volumes:
@@ -43,6 +44,7 @@ services:
             --experimental.localPlugins.captcha-protect.moduleName=github.com/libops/captcha-protect
         volumes:
             - /var/run/docker.sock:/var/run/docker.sock:z
+            - ./tmp:/tmp:rw
             - ./../:/plugins-local/src/github.com/libops/captcha-protect:r
         ports:
             - "80:80"

ci/test.go

@@ -1,7 +1,9 @@
 package main
 
 import (
+	"encoding/json"
 	"fmt"
+	"io"
 	"log"
 	"log/slog"
 	"math/rand"
@@ -54,14 +56,29 @@ func main() {
 	fmt.Printf("Making sure attempt #%d causes a redirect to the challenge page\n", rateLimit+1)
 	ensureRedirect(ips)
 
-	fmt.Println("Sleeping for 3m")
-	time.Sleep(3 * time.Minute)
+	fmt.Println("Sleeping for 2m")
+	time.Sleep(125 * time.Second)
 	fmt.Println("Making sure one attempt passes after 2m window")
 	runParallelChecks(ips, 1)
-
 	fmt.Println("All good 🚀")
 
+	// make sure the state has time to save
+	fmt.Println("Waiting for state to save")
+	runCommand("jq", ".", "tmp/state.json")
+	time.Sleep(80 * time.Second)
+	runCommand("jq", ".", "tmp/state.json")
+
 	runCommand("docker", "container", "stats", "--no-stream")
+
+	// now restart the containers and make sure the previous state reloaded
+	runCommand("docker", "compose", "down")
+	runCommand("docker", "compose", "up", "-d")
+	waitForService("http://localhost")
+	time.Sleep(10 * time.Second)
+	checkStateReload()
+
+	runCommand("rm", "tmp/state.json")
+
 }
 
 func generateUniquePublicIPs(n int) []string {
@@ -191,3 +208,35 @@ func runCommand(name string, args ...string) {
 		log.Fatalf("Command failed: %v", err)
 	}
 }
+
+func checkStateReload() {
+	resp, err := http.Get("http://localhost/captcha-protect/stats")
+	if err != nil {
+		log.Fatalf("Failed to make GET request: %v", err)
+	}
+	defer resp.Body.Close()
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		log.Fatalf("Failed to read response body: %v", err)
+	}
+	var jsonResponse map[string]interface{}
+	err = json.Unmarshal(body, &jsonResponse)
+	if err != nil {
+		log.Fatalf("Failed to unmarshal JSON: %v", err)
+	}
+	bots, exists := jsonResponse["bots"]
+	if !exists {
+		log.Fatalf("Key 'bots' not found in JSON response")
+	}
+	botsMap, ok := bots.(map[string]interface{})
+	if !ok {
+		log.Fatalf("'bots' is not an array")
+	}
+
+	if len(botsMap) != numIPs {
+		log.Fatalf("Expected %d bots, but got %d", numIPs, len(botsMap))
+	}
+
+	log.Println("State reloaded successfully!")
+}

internal/state/state.go

@@ -0,0 +1,49 @@
+package state
+
+import (
+	"reflect"
+
+	lru "github.com/patrickmn/go-cache"
+)
+
+type State struct {
+	Rate     map[string]uint    `json:"rate"`
+	Bots     map[string]bool    `json:"bots"`
+	Verified map[string]bool    `json:"verified"`
+	Memory   map[string]uintptr `json:"memory"`
+}
+
+func GetState(rateCache, botCache, verifiedCache map[string]lru.Item) State {
+	state := State{
+		Memory: make(map[string]uintptr, 3),
+	}
+
+	state.Rate = make(map[string]uint, len(rateCache))
+	state.Memory["rate"] = reflect.TypeOf(state.Rate).Size()
+	for k, v := range rateCache {
+		state.Rate[k] = v.Object.(uint)
+		state.Memory["rate"] += reflect.TypeOf(k).Size()
+		state.Memory["rate"] += reflect.TypeOf(v).Size()
+		state.Memory["rate"] += uintptr(len(k))
+	}
+
+	state.Bots = make(map[string]bool, len(botCache))
+	state.Memory["bot"] = reflect.TypeOf(state.Bots).Size()
+	for k, v := range botCache {
+		state.Bots[k] = v.Object.(bool)
+		state.Memory["bot"] += reflect.TypeOf(k).Size()
+		state.Memory["bot"] += reflect.TypeOf(v).Size()
+		state.Memory["bot"] += uintptr(len(k))
+	}
+
+	state.Verified = make(map[string]bool, len(verifiedCache))
+	state.Memory["verified"] = reflect.TypeOf(state.Verified).Size()
+	for k, v := range verifiedCache {
+		state.Verified[k] = v.Object.(bool)
+		state.Memory["verified"] += reflect.TypeOf(k).Size()
+		state.Memory["verified"] += reflect.TypeOf(v).Size()
+		state.Memory["verified"] += uintptr(len(k))
+	}
+
+	return state
+}