Complete limiter middleware

This commit is contained in:
VnPower
2024-08-28 18:15:21 +07:00
committed by iacore
parent 839a1be6ba
commit 39b5cb8ce0
2 changed files with 25 additions and 4 deletions
+23 -3
View File
@@ -2,12 +2,16 @@ package handlers
import (
"errors"
"math"
"net"
"net/http"
"strings"
"sync"
"golang.org/x/time/rate"
"codeberg.org/vnpower/pixivfe/v2/config"
"codeberg.org/vnpower/pixivfe/v2/routes"
)
func CanRequestSkipLimiter(r *http.Request) bool {
@@ -15,7 +19,8 @@ func CanRequestSkipLimiter(r *http.Request) bool {
return strings.HasPrefix(path, "/img/") ||
strings.HasPrefix(path, "/css/") ||
strings.HasPrefix(path, "/js/") ||
strings.HasPrefix(path, "/proxy/s.pximg.net/")
strings.HasPrefix(path, "/proxy/s.pximg.net/") ||
strings.HasPrefix(path, "/favicon.ico")
}
// Todo: Should we put middlewares in a separate file?
@@ -54,7 +59,15 @@ func (lim *IPRateLimiter) Allow(ip string) bool {
return rl.Allow()
}
var limiter *IPRateLimiter = NewIPRateLimiter(0, 0)
var limiter *IPRateLimiter
func InitializeRateLimiter() {
r := float64(config.GlobalServerConfig.RequestLimit) / 30.0
if config.GlobalServerConfig.RequestLimit < 1 {
r = math.Inf(1)
}
limiter = NewIPRateLimiter(rate.Limit(r), 3)
}
func RateLimitRequest(handler http.HandlerFunc) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
@@ -67,8 +80,15 @@ func RateLimitRequest(handler http.HandlerFunc) http.HandlerFunc {
if !limiter.Allow(ip) {
CatchError(func(w http.ResponseWriter, r *http.Request) error {
err := errors.New("Too many requests")
GetUserContext(r).Err = err
GetUserContext(r).StatusCode = http.StatusTooManyRequests
return errors.New("Too many requests")
err = routes.ErrorPage(w, r, err)
if err != nil {
println("Error rendering error route: %s", err)
}
return err
})(w, r)
} else {
handler(w, r)
+2 -1
View File
@@ -28,6 +28,7 @@ func main() {
ctx_timeout, cancel := context.WithTimeout(context.Background(), config.ProxyCheckerTimeout)
defer cancel()
config.InitializeProxyChecker(ctx_timeout)
handlers.InitializeRateLimiter()
router := handlers.DefineRoutes()
@@ -35,7 +36,7 @@ func main() {
router.ServeHTTP(w, r)
handlers.ErrorHandler(w, r)
}
// main_handler = handlers.RateLimitRequest(main_handler)
main_handler = handlers.RateLimitRequest(main_handler)
main_handler = handlers.LogRequest(main_handler)
// run sass when in development mode