rate limit with burst

bkawk · Feb 4, 2023 · ee10c74 · ee10c74
1 parent a899794
commit ee10c74
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 0 deletions.
diff --git a/api/main.go b/api/main.go
@@ -5,9 +5,11 @@ import (
 	"bkawk/go-echo/api/handlers"
 	"context"
 	"fmt"
+	"net/http"
 	"os"
 
 	"github.com/joho/godotenv"
+	"github.com/juju/ratelimit"
 	"github.com/labstack/echo/v4"
 )
 
@@ -38,6 +40,17 @@ func main() {
 		}
 	})
 
+	// Limit the number of requests to 1 requests per second with a burst of 20 requests
+	limiter := ratelimit.NewBucketWithQuantum(1, 20, 1)
+	e.Use(func(next echo.HandlerFunc) echo.HandlerFunc {
+		return func(c echo.Context) error {
+			if limiter.TakeAvailable(1) == 0 {
+				return c.String(http.StatusTooManyRequests, "Rate limit exceeded")
+			}
+			return next(c)
+		}
+	})
+
 	// Routes
 	e.GET("/health", handlers.HealthGet)                            // Health Check
 	e.POST("/register", handlers.RegisterPost)                      // Register a new user

diff --git a/go.mod b/go.mod
@@ -4,6 +4,7 @@ go 1.19
 
 require (
 	github.com/joho/godotenv v1.5.0
+	github.com/juju/ratelimit v1.0.2
 	github.com/labstack/echo/v4 v4.10.0
 	go.mongodb.org/mongo-driver v1.11.1
 )

diff --git a/go.sum b/go.sum
@@ -7,6 +7,8 @@ github.com/google/go-cmp v0.5.2 h1:X2ev0eStA3AbceY54o37/0PQ/UWqKEiiO2dKL5OPaFM=
 github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/joho/godotenv v1.5.0 h1:C/Vohk/9L1RCoS/UW2gfyi2N0EElSW3yb9zwi3PjosE=
 github.com/joho/godotenv v1.5.0/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
+github.com/juju/ratelimit v1.0.2 h1:sRxmtRiajbvrcLQT7S+JbqU0ntsb9W2yhSdNN8tWfaI=
+github.com/juju/ratelimit v1.0.2/go.mod h1:qapgC/Gy+xNh9UxzV13HGGl/6UXNN+ct+vwSgWNm/qk=
 github.com/klauspost/compress v1.13.6 h1:P76CopJELS0TiO2mebmnzgWaajssP/EszplttgQxcgc=
 github.com/klauspost/compress v1.13.6/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
 github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
@@ -78,6 +80,7 @@ golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGm
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=