From c0139035115363efb2e2598132336524183e0d9f Mon Sep 17 00:00:00 2001
From: Kirill Sysoev <k.sysoev@me.com>
Date: Wed, 10 Apr 2024 19:42:26 +0800
Subject: [PATCH] Adds rate limiter middleware

---
 go.mod                            |  3 ++-
 go.sum                            |  2 ++
 middleware/request/ratelimiter.go | 24 ++++++++++++++++++++++++
 3 files changed, 28 insertions(+), 1 deletion(-)
 create mode 100644 middleware/request/ratelimiter.go

diff --git a/go.mod b/go.mod
index c8aa085..6867533 100644
--- a/go.mod
+++ b/go.mod
@@ -1,6 +1,6 @@
 module github.com/ksysoev/wasabi
 
-go 1.21.4
+go 1.22.1
 
 require (
 	github.com/google/uuid v1.5.0
@@ -11,6 +11,7 @@ require (
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/ksysoev/ratestor v0.1.0 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/stretchr/objx v0.5.2 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
diff --git a/go.sum b/go.sum
index 4774800..10e1689 100644
--- a/go.sum
+++ b/go.sum
@@ -2,6 +2,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/google/uuid v1.5.0 h1:1p67kYwdtXjb0gL0BPiP1Av9wiZPo5A8z2cWkTZ+eyU=
 github.com/google/uuid v1.5.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/ksysoev/ratestor v0.1.0 h1:zAlHYNXHyfwj78TnjUF6FHyYwkMcZxxxGul2DRhF4/c=
+github.com/ksysoev/ratestor v0.1.0/go.mod h1:ZJ3MX2d9JtBetKh9WMLvn0ESotKPJnl9rEU/qetyObk=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
diff --git a/middleware/request/ratelimiter.go b/middleware/request/ratelimiter.go
new file mode 100644
index 0000000..4ccbdda
--- /dev/null
+++ b/middleware/request/ratelimiter.go
@@ -0,0 +1,24 @@
+package request
+
+import (
+	"time"
+
+	"github.com/ksysoev/ratestor"
+	"github.com/ksysoev/wasabi"
+	"github.com/ksysoev/wasabi/dispatch"
+)
+
+func NewRateLimiterMiddleware(requestLimit func(wasabi.Request) (key string, period time.Duration, limit uint64)) func(next wasabi.RequestHandler) wasabi.RequestHandler {
+	stor := ratestor.NewRateStor()
+
+	return func(next wasabi.RequestHandler) wasabi.RequestHandler {
+		return dispatch.RequestHandlerFunc(func(conn wasabi.Connection, req wasabi.Request) error {
+			key, period, limit := requestLimit(req)
+			if err := stor.Allow(key, period, limit); err != nil {
+				return err
+			}
+
+			return next.Handle(conn, req)
+		})
+	}
+}