From d8e9afd8dad5861a064283168bdf4a0e3f18b170 Mon Sep 17 00:00:00 2001
From: Iris Lightshard <nilix@nilfm.cc>
Date: Thu, 28 Nov 2024 09:54:35 -0700
Subject: [PATCH] add rateLimiters and Throttle middleware

---
 README.md                  |  4 +-
 middleware/middleware.go   | 12 ++++++
 quartzgun_test.go          |  9 ++++-
 rateLimiter/rateLimiter.go | 77 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 100 insertions(+), 2 deletions(-)
 create mode 100644 rateLimiter/rateLimiter.go

diff --git a/README.md b/README.md
index 5b20f30..1ae23c9 100644
--- a/README.md
+++ b/README.md
@@ -18,7 +18,7 @@ Thinking about URL routes reminded me of the tree of light the fictional [Quartz
 
 ## usage
 
-A more complete usage guide will be forthcoming, but for now you can check out the [quartzgun_test.go](https://nilfm.cc/git/quartzgun/tree/quartzgun_test.go) file for an overview of how to use it.
+You can check out the [quartzgun_test.go](./quartzgun_test.go) file for an overview of how to use it, or see projects like [nirvash](https://forge.lightcrystal.system/nilix/nirvash) and [felt](https://forge.lightcrystal.systems/nilix/felt) which use quartzgun extensively.
 
 ## roadmap/features
 
@@ -28,6 +28,7 @@ Features may be added here at any time as things are in early stages right now:
 
 * [x] router (static service trees, paramaterized routes, and per-method handlers on routes)
 * [x] basic renderers (HTML template, JSON, XML)
+* [x] rate limiters (one by IP and one that is indiscriminate)
 
 ### auth
 
@@ -45,6 +46,7 @@ Features may be added here at any time as things are in early stages right now:
   - [x] `Defend`: enact CSRF protection (use on the endpoint)
   - [x] `Provision`: use BASIC authentication to provision an access token
   - [x] `Validate`: valiate the bearer token against the `UserStore`
+  - [x] `Throttle`: rate limit using a `func(*http.Request)bool`
 
 ## license
 
diff --git a/middleware/middleware.go b/middleware/middleware.go
index d163f6f..2f5adcb 100644
--- a/middleware/middleware.go
+++ b/middleware/middleware.go
@@ -167,3 +167,15 @@ func Defend(next http.Handler, userStore auth.UserStore, denied string) http.Han
 
 	return http.HandlerFunc(handlerFunc)
 }
+
+func Throttle(next http.Handler, bouncer func(*http.Request) bool) http.Handler {
+	handlerFunc := func(w http.ResponseWriter, req *http.Request) {
+		allowed := bouncer(req)
+		if allowed {
+			next.ServeHTTP(w, req)
+			return
+		}
+		w.WriteHeader(http.StatusTooManyRequests)
+	}
+	return http.HandlerFunc(handlerFunc)
+}
diff --git a/quartzgun_test.go b/quartzgun_test.go
index 5c7251f..d550ba3 100644
--- a/quartzgun_test.go
+++ b/quartzgun_test.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"hacklab.nilfm.cc/quartzgun/indentalUserDB"
 	"hacklab.nilfm.cc/quartzgun/middleware"
+	"hacklab.nilfm.cc/quartzgun/rateLimiter"
 	"hacklab.nilfm.cc/quartzgun/renderer"
 	"hacklab.nilfm.cc/quartzgun/router"
 	"html/template"
@@ -36,6 +37,12 @@ func TestMain(m *testing.M) {
 	udb.AddUser("nilix", "questing")
 	sesh, _ := udb.InitiateSession("nilix", "questing", 60)
 
+	bouncer := rateLimiter.IpRateLimiter{
+		map[string]*rateLimiter.RateLimitData{},
+		5,
+		2,
+	}
+
 	fmt.Printf("%s // %s\n", sesh, sesh)
 	rtr := &router.Router{
 		StaticPaths: map[string]string{
@@ -56,7 +63,7 @@ func TestMain(m *testing.M) {
 		renderer.Template(
 			"testData/templates/test.html"), http.MethodGet, udb, "/login"))
 
-	rtr.Get("/json", ApiSomething(renderer.JSON("apiData")))
+	rtr.Get("/json", middleware.Throttle(ApiSomething(renderer.JSON("apiData")), bouncer.RateLimit))
 
 	rtr.Get(`/thing/(?P<Thing>\w+)`, renderer.Template("testData/templates/paramTest.html"))
 
diff --git a/rateLimiter/rateLimiter.go b/rateLimiter/rateLimiter.go
new file mode 100644
index 0000000..bbb1280
--- /dev/null
+++ b/rateLimiter/rateLimiter.go
@@ -0,0 +1,77 @@
+package rateLimiter
+
+import (
+	"net/http"
+	"strings"
+	"time"
+)
+
+type RateLimitData struct {
+	Attempts   int
+	LastAccess time.Time
+}
+
+// limit by IP
+
+type IpRateLimiter struct {
+	Data            map[string]*RateLimitData
+	Seconds         int
+	AttemptsAllowed int
+}
+
+func (self *IpRateLimiter) RateLimit(req *http.Request) bool {
+	curtime := time.Now()
+
+	addrWithPort := req.RemoteAddr
+	addrParts := strings.Split(addrWithPort, ":")
+	addr := strings.Join(addrParts[:len(addrParts)-1], ":")
+
+	data, exists := self.Data[addr]
+	if !exists {
+		self.Data[addr] = &RateLimitData{1, curtime}
+		return true
+	}
+
+	if curtime.Sub(data.LastAccess).Seconds() > float64(self.Seconds) {
+		data.Attempts = 1
+		data.LastAccess = curtime
+		return true
+	}
+
+	if self.AttemptsAllowed > data.Attempts {
+		data.Attempts++
+		data.LastAccess = curtime
+		return true
+	}
+
+	data.Attempts++
+	return false
+}
+
+// limit regardless of IP
+
+type IndiscriminateRateLimiter struct {
+	Attempts        int
+	AttemptsAllowed int
+	Seconds         int
+	LastAccess      time.Time
+}
+
+func (self *IndiscriminateRateLimiter) RateLimit(req *http.Request) bool {
+	curtime := time.Now()
+
+	if curtime.Sub(self.LastAccess).Seconds() > float64(self.Seconds) {
+		self.Attempts = 1
+		self.LastAccess = curtime
+		return true
+	}
+
+	if self.AttemptsAllowed > self.Attempts {
+		self.Attempts++
+		self.LastAccess = curtime
+		return true
+	}
+
+	self.Attempts++
+	return false
+}