feat: ratelimit middleware

src/cmd/main.go

@@ -44,6 +44,10 @@ func main() {
 		MaxAge:        12 * time.Hour,
 	}))
 
+	ratelimitRepo := repository.NewRatelimitRepository()
+
+	middlewareController := controllers.NewMiddlewareController(ratelimitRepo)
+
 	apiRepo := repository.ApiRepository{
 		DB: db,
 	}
@@ -52,7 +56,7 @@ func main() {
 
 	pasteController := controllers.NewPasteController(pasteRepository)
 
-	r.POST("/new", pasteController.CreatePaste)
+	r.POST("/new", middlewareController.RatelimitMiddleware(), pasteController.CreatePaste)
 	r.GET("/metadata/:pasteID", pasteController.GetPasteMetadata)
 	r.GET("/data/:pasteID", pasteController.GetPaste)
 

src/controllers/middleware.go

@@ -0,0 +1,32 @@
+package controllers
+
+import (
+	"carbon-paper/src/repository"
+
+	"github.com/gin-gonic/gin"
+)
+
+type MiddlewareController struct {
+	RatelimitRepository *repository.RatelimitRepository
+}
+
+func NewMiddlewareController(repo *repository.RatelimitRepository) *MiddlewareController {
+	return &MiddlewareController{
+		RatelimitRepository: repo,
+	}
+}
+
+func (ctrl *MiddlewareController) RatelimitMiddleware() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		if ctrl.RatelimitRepository.Get(c.ClientIP()) > 10 {
+			c.AbortWithStatusJSON(429, gin.H{
+				"message": "Too many requests",
+			})
+			return
+		}
+
+		ctrl.RatelimitRepository.Increment(c.ClientIP())
+
+		c.Next()
+	}
+}

src/repository/ratelimitRepo.go

@@ -0,0 +1,50 @@
+package repository
+
+import (
+	"sync"
+	"time"
+)
+
+type RatelimitRepository struct {
+	ipMap *sync.Map
+}
+
+func NewRatelimitRepository() *RatelimitRepository {
+	ipMap := sync.Map{}
+	return &RatelimitRepository{
+		ipMap: &ipMap,
+	}
+}
+
+func (repo *RatelimitRepository) Increment(ip string) {
+	val, _ := repo.ipMap.LoadOrStore(ip, 1)
+	repo.ipMap.Store(ip, val.(int)+1)
+
+	time.AfterFunc(time.Minute, func() {
+		repo.decrement(ip)
+	})
+}
+
+func (repo *RatelimitRepository) decrement(ip string) {
+	val, ok := repo.ipMap.Load(ip)
+	if !ok {
+		return
+	}
+	repo.ipMap.Store(ip, val.(int)-1)
+
+	if val.(int) <= 0 {
+		repo.ipMap.Delete(ip)
+	}
+}
+
+func (repo *RatelimitRepository) Get(ip string) int {
+	val, ok := repo.ipMap.Load(ip)
+	if !ok {
+		return 0
+	}
+	return val.(int)
+}
+
+func (repo *RatelimitRepository) Delete(ip string) {
+	repo.ipMap.Delete(ip)
+}