From 39b5cb8ce033161876d3a03452a8bc3fdc71f7eb Mon Sep 17 00:00:00 2001
From: VnPower <vnpower@disroot.org>
Date: Wed, 28 Aug 2024 18:15:21 +0700
Subject: [PATCH] Complete limiter middleware

---
 handlers/limiter.go | 26 +++++++++++++++++++++++---
 main.go             |  3 ++-
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/handlers/limiter.go b/handlers/limiter.go
index 5cac332..13fdc73 100644
--- a/handlers/limiter.go
+++ b/handlers/limiter.go
@@ -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)
diff --git a/main.go b/main.go
index 1772d94..dc9d717 100644
--- a/main.go
+++ b/main.go
@@ -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