diff --git a/main.go b/main.go index e42615e..0a3f775 100644 --- a/main.go +++ b/main.go @@ -14,7 +14,9 @@ func main() { r.LoadHTMLGlob("templates/*") r.Static("/static", "./public") + r.Use(gin.Recovery()) r.Use(middleware.OptionsMiddleware()) + r.Use(middleware.Ratelimit()) r.GET("/robots.txt", func(c *gin.Context) { c.String(200, "User-agent: *\nDisallow: /") diff --git a/src/middleware/ratelimit.go b/src/middleware/ratelimit.go new file mode 100644 index 0000000..e386513 --- /dev/null +++ b/src/middleware/ratelimit.go @@ -0,0 +1,44 @@ +package middleware + +import ( + "sync" + "time" + + "github.com/gin-gonic/gin" +) + +var ipMap = sync.Map{} + +func Ratelimit() gin.HandlerFunc { + return func(c *gin.Context) { + ip := c.ClientIP() + + // log request count as the value, ip as key + // if the ip is not in the map, create a new entry with a value of 1 + // if they exceed 30 requests in 1 minute, return a 429 + + // get the value from the map + val, ok := ipMap.Load(ip) + if !ok { + // if the ip is not in the map, create a new entry with a value of 1 + ipMap.Store(ip, 1) + c.Next() + return + } + + // if the ip is in the map, increment the value + ipMap.Store(ip, val.(int)+1) + + // if they exceed 30 requests in 1 minute, return a 429 + if val.(int) > 30 { + c.String(429, "429 Too Many Requests") + c.Abort() + return + } + + // delete the ip from the map after 1 minute + time.AfterFunc(time.Minute, func() { + ipMap.Delete(ip) + }) + } +}