Merge pull request #46 from COS301-SE-2024/feat/backend/rate-limit-api

Feat/backend/rate limit api
COS301-SE-2024 · Jun 6, 2024 · 01cff59 · 01cff59
2 parents 358fd0f + 78692b5
commit 01cff59
Show file tree

Hide file tree

Showing 9 changed files with 232 additions and 13 deletions.
diff --git a/.github/workflows/lint-test-build-golang.yml b/.github/workflows/lint-test-build-golang.yml
@@ -54,7 +54,7 @@ jobs:
 
     - name: Run tests
       run: |
-        go test ./tests/... -race -coverprofile=coverage.out -covermode=atomic
+        go test -v -coverpkg=github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils,github.com/COS301-SE-2024/occupi/occupi-backend/pkg/handlers ./tests/... -coverprofile=coverage.out
 
     - name: Upload coverage reports to Codecov
       uses: codecov/[email protected]

diff --git a/occupi-backend/cmd/occupi-backend/main.go b/occupi-backend/cmd/occupi-backend/main.go
@@ -10,6 +10,7 @@ import (
 
 	"github.com/COS301-SE-2024/occupi/occupi-backend/configs"
 	"github.com/COS301-SE-2024/occupi/occupi-backend/pkg/database"
+	"github.com/COS301-SE-2024/occupi/occupi-backend/pkg/middleware"
 	"github.com/COS301-SE-2024/occupi/occupi-backend/pkg/router"
 	"github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils"
 )
@@ -39,6 +40,9 @@ func main() {
 		logrus.Fatal("Failed to set trusted proxies: ", err)
 	}
 
+	// adding rate limiting middleware
+	middleware.AttachRateLimitMiddleware(ginRouter)
+
 	// Register routes
 	router.OccupiRouter(ginRouter, db)
 

diff --git a/occupi-backend/coverage.out b/occupi-backend/coverage.out
@@ -0,0 +1,47 @@
+mode: set
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/OTP.go:9.36,12.16 3 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/OTP.go:12.16,14.3 1 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/OTP.go:15.2,16.21 2 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/response.go:16.9,22.2 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/response.go:30.9,37.2 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/response.go:46.9,56.2 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/response.go:59.34,69.2 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/response.go:79.9,90.2 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:22.20,27.16 3 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:27.16,29.3 1 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:33.2,42.35 3 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:46.42,50.16 4 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:50.16,52.3 1 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:53.2,55.17 3 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:59.34,62.16 3 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:62.16,64.3 1 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:65.2,66.19 2 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:70.44,73.16 3 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:73.16,75.3 1 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:77.2,79.19 2 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:83.94,89.31 4 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:89.31,91.25 2 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:91.25,93.62 2 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:93.62,97.5 3 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:102.2,104.20 2 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:108.41,111.2 2 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:114.39,118.2 2 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:121.45,130.23 2 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:130.23,132.3 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:134.2,134.44 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:134.44,136.3 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:138.2,138.44 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:138.44,140.3 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:142.2,142.34 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:142.34,144.3 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:146.2,146.40 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:146.40,148.3 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:150.2,150.13 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:154.35,158.2 2 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:161.52,167.16 2 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:167.16,169.3 1 0
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:170.2,170.18 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:174.80,178.16 2 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:178.16,180.3 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:181.2,181.19 1 1
+github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils/utils.go:184.20,187.2 0 0
diff --git a/occupi-backend/go.mod b/occupi-backend/go.mod
@@ -47,11 +47,13 @@ require (
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/montanaflynn/stats v0.7.1 // indirect
 	github.com/pelletier/go-toml/v2 v2.2.2 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/rogpeppe/go-internal v1.9.0 // indirect
 	github.com/stretchr/objx v0.5.2 // indirect
 	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
 	github.com/ugorji/go/codec v1.2.12 // indirect
+	github.com/ulule/limiter/v3 v3.11.2 // indirect
 	github.com/xdg-go/pbkdf2 v1.0.0 // indirect
 	github.com/xdg-go/scram v1.1.2 // indirect
 	github.com/xdg-go/stringprep v1.0.4 // indirect

diff --git a/occupi-backend/go.sum b/occupi-backend/go.sum
@@ -84,6 +84,8 @@ github.com/montanaflynn/stats v0.7.1 h1:etflOAAHORrCC44V+aR6Ftzort912ZU+YLiSTuV8
 github.com/montanaflynn/stats v0.7.1/go.mod h1:etXPPgVO6n31NxCd9KQUMvCM+ve0ruNzt6R8Bnaayow=
 github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM=
 github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs=
+github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
+github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 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/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc=
@@ -108,6 +110,8 @@ github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS
 github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
 github.com/ugorji/go/codec v1.2.12 h1:9LC83zGrHhuUA9l16C9AHXAqEV/2wBQ4nkvumAE65EE=
 github.com/ugorji/go/codec v1.2.12/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg=
+github.com/ulule/limiter/v3 v3.11.2 h1:P4yOrxoEMJbOTfRJR2OzjL90oflzYPPmWg+dvwN2tHA=
+github.com/ulule/limiter/v3 v3.11.2/go.mod h1:QG5GnFOCV+k7lrL5Y8kgEeeflPH3+Cviqlqa8SVSQxI=
 github.com/xdg-go/pbkdf2 v1.0.0 h1:Su7DPu48wXMwC3bs7MCNG+z4FhcyEuz5dlvchbq0B0c=
 github.com/xdg-go/pbkdf2 v1.0.0/go.mod h1:jrpuAogTd400dnrH08LKmI/xc1MbPOebTwRqcT5RDeI=
 github.com/xdg-go/scram v1.1.2 h1:FHX5I5B4i4hKRVRBCFRxq1iQRej7WO3hhBuJf+UUySY=

diff --git a/occupi-backend/occupi.bat b/occupi-backend/occupi.bat
@@ -20,10 +20,10 @@ if "%1 %2" == "run dev" (
     docker-compose up
     exit /b 0
 ) else if "%1" == "test" (
-    go test ./tests/...
+    go test -v ./tests/...
     exit /b 0
 ) else if "%1 %2" == "test codecov" (
-    go test ./tests/... -race -coverprofile=coverage.out -covermode=atomic
+    go test -v -coverpkg=github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils,github.com/COS301-SE-2024/occupi/occupi-backend/pkg/handlers ./tests/... -coverprofile=coverage.out
     exit /b 0
 ) else if "%1" == "lint" (
     golangci-lint run

diff --git a/occupi-backend/occupi.sh b/occupi-backend/occupi.sh
@@ -29,9 +29,9 @@ elif [ "$1" = "docker" ] && [ "$2" = "build" ]; then
 elif [ "$1" = "docker" ] && [ "$2" = "up" ]; then
     docker-compose up
 elif [ "$1" = "test" ]; then
-    go test ./tests/...
+    go test -v ./tests/...
 elif [ "$1" = "test" ] && [ "$2" = "codecov" ]; then
-    go test ./tests/... -race -coverprofile=coverage.out -covermode=atomic
+    go test -v -coverpkg=github.com/COS301-SE-2024/occupi/occupi-backend/pkg/utils,github.com/COS301-SE-2024/occupi/occupi-backend/pkg/handlers ./tests/... -coverprofile=coverage.out
 elif [ "$1" = "lint" ]; then
     golangci-lint run
 elif [ "$1" = "help" ] || [ -z "$1" ]; then

diff --git a/occupi-backend/pkg/middleware/middleware.go b/occupi-backend/pkg/middleware/middleware.go
@@ -5,34 +5,52 @@ import (
 
 	"github.com/gin-contrib/sessions"
 	"github.com/gin-gonic/gin"
+	"github.com/ulule/limiter/v3"
+	mgin "github.com/ulule/limiter/v3/drivers/middleware/gin"
+	"github.com/ulule/limiter/v3/drivers/store/memory"
 )
 
 // ProtectedRoute is a middleware that checks if
 // the user has already been authenticated previously.
-func ProtectedRoute(c *gin.Context) {
-	if sessions.Default(c).Get("profile") == nil {
+func ProtectedRoute(ctx *gin.Context) {
+	if sessions.Default(ctx).Get("profile") == nil {
 		// If the user is not authenticated, return a 401 Unauthorized response
-		c.JSON(http.StatusUnauthorized, gin.H{
+		ctx.JSON(http.StatusUnauthorized, gin.H{
 			"status":  http.StatusUnauthorized,
 			"message": "Bad Request",
 			"error":   "User not authenticated",
 		})
 	} else {
-		c.Next()
+		ctx.Next()
 	}
 }
 
 // ProtectedRoute is a middleware that checks if
 // the user has not been authenticated previously.
-func UnProtectedRoute(c *gin.Context) {
-	if sessions.Default(c).Get("profile") != nil {
+func UnProtectedRoute(ctx *gin.Context) {
+	if sessions.Default(ctx).Get("profile") != nil {
 		// If the user is authenticated, return a 401 Unauthorized response
-		c.JSON(http.StatusUnauthorized, gin.H{
+		ctx.JSON(http.StatusUnauthorized, gin.H{
 			"status":  http.StatusUnauthorized,
 			"message": "Bad Request",
 			"error":   "User already authenticated",
 		})
 	} else {
-		c.Next()
+		ctx.Next()
 	}
 }
+
+// AttachRateLimitMiddleware attaches the rate limit middleware to the router.
+func AttachRateLimitMiddleware(ginRouter *gin.Engine) {
+	// Define a rate limit: 5 requests per second
+	rate, _ := limiter.NewRateFromFormatted("5-S")
+
+	store := memory.NewStore()
+	instance := limiter.New(store, rate)
+
+	// Create the rate limiting middleware
+	middleware := mgin.NewMiddleware(instance)
+
+	// Apply the middleware to the router
+	ginRouter.Use(middleware)
+}
diff --git a/occupi-backend/tests/handlers_test.go b/occupi-backend/tests/handlers_test.go
@@ -4,9 +4,15 @@ import (
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
+	"sync"
 	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
 
 	"github.com/gin-gonic/gin"
+
+	"github.com/COS301-SE-2024/occupi/occupi-backend/pkg/middleware"
 	// "github.com/joho/godotenv"
 	// "github.com/stretchr/testify/assert"
 	// "github.com/stretchr/testify/mock"
@@ -175,6 +181,144 @@ func TestPingRoute(t *testing.T) {
 	}
 }
 
+func TestRateLimit(t *testing.T) {
+	// Create a new Gin router
+	router := gin.Default()
+
+	// attach rate limit middleware
+	middleware.AttachRateLimitMiddleware(router)
+
+	// Register the route
+	router.GET("/ping", func(c *gin.Context) {
+		c.JSON(http.StatusOK, gin.H{"message": "pong -> I am alive and kicking"})
+	})
+
+	server := httptest.NewServer(router)
+	defer server.Close()
+
+	var wg sync.WaitGroup
+	numRequests := 10
+	responseCodes := make([]int, numRequests)
+
+	for i := 0; i < numRequests; i++ {
+		wg.Add(1)
+		go func(index int) {
+			defer wg.Done()
+			resp, err := http.Get(server.URL + "/ping")
+			if err != nil {
+				t.Errorf("Request %d failed: %v", index, err)
+				return
+			}
+			defer resp.Body.Close()
+			responseCodes[index] = resp.StatusCode
+		}(i)
+		time.Sleep(100 * time.Millisecond) // Slight delay to spread out the requests
+	}
+
+	wg.Wait()
+
+	rateLimitedCount := 0
+	for _, code := range responseCodes {
+		if code == http.StatusTooManyRequests {
+			rateLimitedCount++
+		}
+	}
+
+	assert.Greater(t, rateLimitedCount, 0, "There should be some requests that are rate limited")
+	assert.LessOrEqual(t, rateLimitedCount, numRequests-5, "There should be at least 5 requests that are not rate limited")
+}
+
+func TestRateLimitWithMultipleIPs(t *testing.T) {
+	// Create a new Gin router
+	router := gin.Default()
+
+	// attach rate limit middleware
+	middleware.AttachRateLimitMiddleware(router)
+
+	// Register the route
+	router.GET("/ping", func(c *gin.Context) {
+		c.JSON(http.StatusOK, gin.H{"message": "pong -> I am alive and kicking"})
+	})
+
+	server := httptest.NewServer(router)
+	defer server.Close()
+
+	var wg sync.WaitGroup
+	numRequests := 10
+	ip1 := "192.168.1.1"
+	ip2 := "192.168.1.2"
+	responseCodesIP1 := make([]int, numRequests)
+	responseCodesIP2 := make([]int, numRequests-5)
+
+	// Send requests from the first IP address
+	for i := 0; i < numRequests; i++ {
+		wg.Add(1)
+		go func(index int) {
+			defer wg.Done()
+			client := &http.Client{}
+			req, err := http.NewRequest("GET", server.URL+"/ping", nil)
+			if err != nil {
+				t.Errorf("Failed to create request: %v", err)
+				return
+			}
+			req.Header.Set("X-Forwarded-For", ip1)
+			resp, err := client.Do(req)
+			if err != nil {
+				t.Errorf("Request failed: %v", err)
+				return
+			}
+			defer resp.Body.Close()
+			responseCodesIP1[index] = resp.StatusCode
+		}(i)
+		time.Sleep(10 * time.Millisecond) // Slight delay to spread out the requests
+	}
+
+	// Send requests from the second IP address
+	for i := 0; i < numRequests-5; i++ {
+		wg.Add(1)
+		go func(index int) {
+			defer wg.Done()
+			client := &http.Client{}
+			req, err := http.NewRequest("GET", server.URL+"/ping", nil)
+			if err != nil {
+				t.Errorf("Failed to create request: %v", err)
+				return
+			}
+			req.Header.Set("X-Forwarded-For", ip2)
+			resp, err := client.Do(req)
+			if err != nil {
+				t.Errorf("Request failed: %v", err)
+				return
+			}
+			defer resp.Body.Close()
+			responseCodesIP2[index] = resp.StatusCode
+		}(i)
+		time.Sleep(10 * time.Millisecond) // Slight delay to spread out the requests
+	}
+
+	wg.Wait()
+
+	rateLimitedCountIP1 := 0
+	rateLimitedCountIP2 := 0
+	for _, code := range responseCodesIP1 {
+		if code == http.StatusTooManyRequests {
+			rateLimitedCountIP1++
+		}
+	}
+	for _, code := range responseCodesIP2 {
+		if code == http.StatusTooManyRequests {
+			rateLimitedCountIP2++
+		}
+	}
+
+	// Assertions for IP1
+	assert.Greater(t, rateLimitedCountIP1, 0, "There should be some requests from IP1 that are rate limited")
+	assert.LessOrEqual(t, rateLimitedCountIP1, numRequests-5, "There should be at least 5 requests from IP1 that are not rate limited")
+
+	// Assertions for IP2
+	assert.Equal(t, rateLimitedCountIP2, 0, "There should be no requests from IP2 that are rate limited")
+}
+
 /*
 func TestGetResource(t *testing.T) {
 	// Load environment variables from .env file