From 6c709b97a19ca16963161ae95fa92539a1682251 Mon Sep 17 00:00:00 2001
From: grumbulon <grumbulon@grumbulon.xyz>
Date: Fri, 20 Jan 2023 18:29:33 -0500
Subject: [PATCH] add rate limiting

---
 go.mod              |  2 ++
 go.sum              |  4 ++++
 internal/api/api.go | 41 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+)

diff --git a/go.mod b/go.mod
index 63c98c7..d9192c0 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,7 @@ go 1.19
 require (
 	github.com/glebarez/sqlite v1.6.0
 	github.com/go-chi/chi/v5 v5.0.8
+	github.com/go-chi/httprate v0.7.1
 	github.com/go-chi/jwtauth/v5 v5.1.0
 	github.com/miekg/dns v1.1.50
 	github.com/swaggo/swag v1.8.9
@@ -14,6 +15,7 @@ require (
 
 require (
 	github.com/KyleBanks/depth v1.2.1 // indirect
+	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/go-openapi/jsonpointer v0.19.5 // indirect
 	github.com/go-openapi/jsonreference v0.20.0 // indirect
 	github.com/go-openapi/spec v0.20.6 // indirect
diff --git a/go.sum b/go.sum
index a62485a..fc28eb2 100644
--- a/go.sum
+++ b/go.sum
@@ -2,6 +2,8 @@ github.com/KyleBanks/depth v1.2.1 h1:5h8fQADFrWtarTdtDudMmGsC7GPbOAu6RVB3ffsVFHc
 github.com/KyleBanks/depth v1.2.1/go.mod h1:jzSb9d0L43HxTQfT+oSA1EEp2q+ne2uh6XgeJcm8brE=
 github.com/adrg/xdg v0.4.0 h1:RzRqFcjH4nE5C6oTAxhBtoE2IRyjBSa62SCbyPidvls=
 github.com/adrg/xdg v0.4.0/go.mod h1:N6ag73EX4wyxeaoeHctc1mas01KZgsj5tYiAIwqJE/E=
+github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
+github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/chzyer/logex v1.2.0/go.mod h1:9+9sk7u7pGNWYMkh0hdiL++6OeibzJccyQU4p4MedaY=
 github.com/chzyer/readline v1.5.0/go.mod h1:x22KAscuvRqlLoK9CsoYsmxoXZMMFVyOl86cAH8qUic=
 github.com/chzyer/test v0.0.0-20210722231415-061457976a23/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU=
@@ -23,6 +25,8 @@ github.com/go-chi/chi/v5 v5.0.8 h1:lD+NLqFcAi1ovnVZpsnObHGW4xb4J8lNmoYVfECH1Y0=
 github.com/go-chi/chi/v5 v5.0.8/go.mod h1:DslCQbL2OYiznFReuXYUmQ2hGd1aDpCnlMNITLSKoi8=
 github.com/go-chi/httplog v0.2.5 h1:S02eG9NTrB/9kk3Q3RA3F6CR2b+v8WzB8IxK+zq3dBo=
 github.com/go-chi/httplog v0.2.5/go.mod h1:/pIXuFSrOdc5heKIJRA5Q2mW7cZCI2RySqFZNFoZjKg=
+github.com/go-chi/httprate v0.7.1 h1:d5kXARdms2PREQfU4pHvq44S6hJ1hPu4OXLeBKmCKWs=
+github.com/go-chi/httprate v0.7.1/go.mod h1:6GOYBSwnpra4CQfAKXu8sQZg+nZ0M1g9QnyFvxrAB8A=
 github.com/go-chi/jwtauth/v5 v5.1.0 h1:wJyf2YZ/ohPvNJBwPOzZaQbyzwgMZZceE1m8FOzXLeA=
 github.com/go-chi/jwtauth/v5 v5.1.0/go.mod h1:MA93hc1au3tAQwCKry+fI4LqJ5MIVN4XSsglOo+lSc8=
 github.com/go-openapi/jsonpointer v0.19.3/go.mod h1:Pl9vOtqEWErmShwVjC8pYs9cog34VGT37dQOVbmoatg=
diff --git a/internal/api/api.go b/internal/api/api.go
index afa57bd..d1068f6 100644
--- a/internal/api/api.go
+++ b/internal/api/api.go
@@ -2,13 +2,16 @@ package api
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"net/http"
 	"time"
 
+	"git.freecumextremist.com/grumbulon/pomme/internal"
 	"git.freecumextremist.com/grumbulon/pomme/internal/db"
 	"github.com/go-chi/chi/v5"
 	"github.com/go-chi/httplog"
+	"github.com/go-chi/httprate"
 	"github.com/go-chi/jwtauth/v5"
 )
 
@@ -56,6 +59,25 @@ func API() http.Handler {
 
 	// Protected routes
 	api.Group(func(api chi.Router) {
+		api.Use(httprate.Limit(
+			10,             // requests
+			10*time.Second, // per duration
+			httprate.WithKeyFuncs(httprate.KeyByIP, httprate.KeyByEndpoint),
+			httprate.WithLimitHandler(func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Set("Content-Type", "application/json")
+				w.WriteHeader(http.StatusTooManyRequests)
+				err := json.NewEncoder(w).Encode(
+					internal.Response{
+						HTTPResponse: http.StatusTooManyRequests,
+						Message:      "API rate limit exceded",
+					})
+				if err != nil {
+					internalServerError(w, "internal server error")
+
+					return
+				}
+			}),
+		))
 		api.Use(jwtauth.Verifier(tokenAuth))
 
 		api.Use(jwtauth.Authenticator)
@@ -65,6 +87,25 @@ func API() http.Handler {
 
 	// Open routes
 	api.Group(func(api chi.Router) {
+		api.Use(httprate.Limit(
+			5,             // requests
+			5*time.Second, // per duration
+			httprate.WithKeyFuncs(httprate.KeyByIP, httprate.KeyByEndpoint),
+			httprate.WithLimitHandler(func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Set("Content-Type", "application/json")
+				w.WriteHeader(http.StatusTooManyRequests)
+				err := json.NewEncoder(w).Encode(
+					internal.Response{
+						HTTPResponse: http.StatusTooManyRequests,
+						Message:      "API rate limit exceded",
+					})
+				if err != nil {
+					internalServerError(w, "internal server error")
+
+					return
+				}
+			}),
+		))
 		api.Use(setDBMiddleware)
 		api.With(setDBMiddleware).Post("/create", NewUser)
 		api.With(setDBMiddleware).Post("/login", Login)