48 lines
1.3 KiB
Go
48 lines
1.3 KiB
Go
|
package main
|
||
|
|
|
||
|
|
import (
|
||
|
|
"log"
|
||
|
|
"net/http"
|
||
|
|
"sync/atomic"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
ratelimit "github.com/JGLTechnologies/gin-rate-limit"
|
||
|
|
"github.com/gin-gonic/gin"
|
||
|
|
)
|
||
|
|
|
||
|
|
func mgRateLimit() gin.HandlerFunc {
|
||
|
|
var totalLimited uint32
|
||
|
|
store := ratelimit.InMemoryStore(&ratelimit.InMemoryOptions{
|
||
|
|
Rate: time.Second,
|
||
|
|
Limit: 50, // TODO: This must be changed to the 1000rps.
|
||
|
|
})
|
||
|
|
return ratelimit.RateLimiter(store, &ratelimit.Options{
|
||
|
|
ErrorHandler: func(c *gin.Context, info ratelimit.Info) {
|
||
|
|
cur := atomic.AddUint32(&totalLimited, 1)
|
||
|
|
log.Printf("%s: [%d] rate limited > %t, will reset at %s\n", c.ClientIP(), cur, info.RateLimited, info.ResetTime)
|
||
|
|
c.AbortWithStatus(http.StatusTooManyRequests)
|
||
|
|
},
|
||
|
|
KeyFunc: keyFunc,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
func vkRateLimit() gin.HandlerFunc {
|
||
|
|
var totalLimited uint32
|
||
|
|
store := ratelimit.InMemoryStore(&ratelimit.InMemoryOptions{
|
||
|
|
Rate: time.Second,
|
||
|
|
Limit: 20,
|
||
|
|
})
|
||
|
|
return ratelimit.RateLimiter(store, &ratelimit.Options{
|
||
|
|
ErrorHandler: func(c *gin.Context, info ratelimit.Info) {
|
||
|
|
cur := atomic.AddUint32(&totalLimited, 1)
|
||
|
|
log.Printf("%s: [%d] VK request rate limited > %t, will reset at %s\n", c.ClientIP(), cur, info.RateLimited, info.ResetTime)
|
||
|
|
c.AbortWithStatus(http.StatusTooManyRequests)
|
||
|
|
},
|
||
|
|
KeyFunc: keyFunc,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
func keyFunc(c *gin.Context) string {
|
||
|
|
return c.ClientIP()
|
||
|
|
}
|