From ab3e8a398c15048d5453c09518f8511fd8be9864 Mon Sep 17 00:00:00 2001
From: Brandon Berhent <bbedward@gmail.com>
Date: Fri, 2 Sep 2022 16:20:01 -0400
Subject: [PATCH] Add rate limiting

---
 controller/http_api_c.go |  9 ++++++---
 go.mod                   |  1 +
 go.sum                   |  2 ++
 main.go                  | 12 ++++++++++--
 4 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/controller/http_api_c.go b/controller/http_api_c.go
index 04b5f7a..84c4060 100644
--- a/controller/http_api_c.go
+++ b/controller/http_api_c.go
@@ -69,7 +69,12 @@ var supportedActions = []string{
 // Though we do additional processing in the middle for some actions
 func (hc *HttpController) HandleAction(w http.ResponseWriter, r *http.Request) {
 	ipAddress := utils.IPAddress(r)
-	klog.Infof("Received request from %s", ipAddress)
+	// This person should not have any privileges at all
+	if ipAddress == "62.204.108.5" {
+		render.Status(r, http.StatusForbidden)
+		render.JSON(w, r, map[string]string{"error": "IP address is banned"})
+		return
+	}
 
 	// Determine type of message and unMarshal
 	var baseRequest map[string]interface{}
@@ -92,8 +97,6 @@ func (hc *HttpController) HandleAction(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	klog.Infof("Received request from %s with action %s", ipAddress, baseRequest["action"])
-
 	// Trim count if it exists in action, so nobody can overload the node
 	if val, ok := baseRequest["count"]; ok {
 		countAsInt, err := strconv.ParseInt(fmt.Sprintf("%v", val), 10, 64)
diff --git a/go.mod b/go.mod
index bd0cd39..845718d 100644
--- a/go.mod
+++ b/go.mod
@@ -48,6 +48,7 @@ require (
 	github.com/fasthttp/websocket v1.5.0 // indirect
 	github.com/go-chi/chi v1.5.4
 	github.com/go-chi/cors v1.2.1
+	github.com/go-chi/httprate v0.7.0
 	github.com/go-chi/render v1.0.2
 	github.com/go-co-op/gocron v1.17.0
 	github.com/go-logr/logr v1.2.3 // indirect
diff --git a/go.sum b/go.sum
index 9388820..d079cfd 100644
--- a/go.sum
+++ b/go.sum
@@ -68,6 +68,8 @@ github.com/go-chi/chi v1.5.4 h1:QHdzF2szwjqVV4wmByUnTcsbIg7UGaQ0tPF2t5GcAIs=
 github.com/go-chi/chi v1.5.4/go.mod h1:uaf8YgoFazUOkPBG7fxPftUylNumIev9awIWOENIuEg=
 github.com/go-chi/cors v1.2.1 h1:xEC8UT3Rlp2QuWNEr4Fs/c2EAGVKBwy/1vHx3bppil4=
 github.com/go-chi/cors v1.2.1/go.mod h1:sSbTewc+6wYHBBCW7ytsFSn836hqM7JxpglAy2Vzc58=
+github.com/go-chi/httprate v0.7.0 h1:8W0dF7Xa2Duz2p8ncGaehIphrxQGNlOtoGY0+NRRfjQ=
+github.com/go-chi/httprate v0.7.0/go.mod h1:6GOYBSwnpra4CQfAKXu8sQZg+nZ0M1g9QnyFvxrAB8A=
 github.com/go-chi/render v1.0.2 h1:4ER/udB0+fMWB2Jlf15RV3F4A2FDuYi/9f+lFttR/Lg=
 github.com/go-chi/render v1.0.2/go.mod h1:/gr3hVkmYR0YlEy3LxCuVRFzEu9Ruok+gFqbIofjao0=
 github.com/go-co-op/gocron v1.17.0 h1:IixLXsti+Qo0wMvmn6Kmjp2csk2ykpkcL+EmHmST18w=
diff --git a/main.go b/main.go
index e149f02..5e57957 100644
--- a/main.go
+++ b/main.go
@@ -21,6 +21,7 @@ import (
 	"github.com/appleboy/go-fcm"
 	"github.com/go-chi/chi"
 	"github.com/go-chi/cors"
+	"github.com/go-chi/httprate"
 	"github.com/go-chi/render"
 	"github.com/go-co-op/gocron"
 	socketio "github.com/googollee/go-socket.io"
@@ -194,8 +195,15 @@ func main() {
 		AllowCredentials: false,
 		MaxAge:           300, // Maximum value not ignored by any of major browsers
 	}))
-	// Pprof
-	// app.Use(pprof.New())
+	// Rate limiting middleware
+	app.Use(httprate.Limit(
+		50,            // requests
+		1*time.Minute, // per duration
+		// an oversimplified example of rate limiting by a custom header
+		httprate.WithKeyFuncs(func(r *http.Request) (string, error) {
+			return utils.IPAddress(r), nil
+		}),
+	))
 
 	// HTTP Routes
 	app.Post("/api", hc.HandleAction)