Use proper rate limiter

Rename package handlers to middleware, since most are middleware

router.go is not middleware. can be moved out to router/router.go

config/config.go

@@ -53,7 +53,7 @@ type ServerConfig struct {
 	APIMaxBackoffTime time.Duration `env:"PIXIVFE_API_MAX_BACKOFF_TIME,overwrite"`
 
 	AcceptLanguage string `env:"PIXIVFE_ACCEPTLANGUAGE,overwrite"`
-	RequestLimit   int    `env:"PIXIVFE_REQUESTLIMIT"` // if 0, request limit is disabled
+	RequestLimit   uint64 `env:"PIXIVFE_REQUESTLIMIT"` // if 0, request limit is disabled
 
 	ProxyServer_staging string  `env:"PIXIVFE_IMAGEPROXY,overwrite"`
 	ProxyServer         url.URL // proxy server URL, may or may not contain authority part of the URL

doc/hosting/environment-variables.md

@@ -62,6 +62,8 @@ The URL of the PixivFE source code repository. This is used in the about page to
 
 **Required**: No
 
+Request limit per half-minute.
+
 Set to a number to enable the built-in rate limiter, e.g., `PIXIVFE_REQUESTLIMIT=15`.
 
 It's recommended to enable rate limiting in the reverse proxy in front of PixivFE rather than using this.

go.mod

@@ -13,10 +13,10 @@ require (
 	github.com/oklog/ulid/v2 v2.1.0
 	github.com/playwright-community/playwright-go v0.4501.1
 	github.com/sethvargo/go-envconfig v1.1.0
+	github.com/sethvargo/go-limiter v1.0.0
 	github.com/soluble-ai/go-jnode v0.1.11
 	github.com/tidwall/gjson v1.17.3
 	golang.org/x/net v0.28.0
-	golang.org/x/time v0.6.0
 )
 
 require (

go.sum

@@ -47,6 +47,8 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/sethvargo/go-envconfig v1.1.0 h1:cWZiJxeTm7AlCvzGXrEXaSTCNgip5oJepekh/BOQuog=
 github.com/sethvargo/go-envconfig v1.1.0/go.mod h1:JLd0KFWQYzyENqnEPWWZ49i4vzZo/6nRidxI8YvGiHw=
+github.com/sethvargo/go-limiter v1.0.0 h1:JqW13eWEMn0VFv86OKn8wiYJY/m250WoXdrjRV0kLe4=
+github.com/sethvargo/go-limiter v1.0.0/go.mod h1:01b6tW25Ap+MeLYBuD4aHunMrJoNO5PVUFdS9rac3II=
 github.com/soluble-ai/go-jnode v0.1.11 h1:To9h6zWI6qr5Vphz0XO4JtCRlkhfilH2EDNwtAZkZT4=
 github.com/soluble-ai/go-jnode v0.1.11/go.mod h1:mKFUM7xxXMBVOygZeE8od2Qls1f3E3NHY1ZhMlzwx/U=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
@@ -106,8 +108,6 @@ golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 golang.org/x/text v0.17.0 h1:XtiM5bkSOt+ewxlOE/aE/AKEHibwj/6gvWMl9Rsh0Qc=
 golang.org/x/text v0.17.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY=
-golang.org/x/time v0.6.0 h1:eTDhh4ZXt5Qf0augr54TN6suAUudPcawVZeIAPU7D4U=
-golang.org/x/time v0.6.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.0.0-20191203134012-c197fd4bf371/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=

main.go

@@ -12,7 +12,7 @@ import (
 
 	"codeberg.org/vnpower/pixivfe/v2/config"
 	"codeberg.org/vnpower/pixivfe/v2/server/audit"
-	"codeberg.org/vnpower/pixivfe/v2/server/handlers"
+	"codeberg.org/vnpower/pixivfe/v2/server/middleware"
 	"codeberg.org/vnpower/pixivfe/v2/server/proxy_checker"
 	"codeberg.org/vnpower/pixivfe/v2/server/template"
 )
@@ -40,15 +40,13 @@ func main() {
 
 	log.Println("Starting server...")
 
-	handlers.InitializeRateLimiter()
-
-	router := handlers.DefineRoutes()
+	router := middleware.DefineRoutes()
 	// the first middleware is the most outer / first executed one
-	router.Use(handlers.ProvideUserContext) // needed for everything else
-	router.Use(handlers.LogRequest)         // all pages need this
-	router.Use(handlers.SetPrivacyHeaders)  // all pages need this
-	router.Use(handlers.HandleError)        // if the inner handler fails, this shows the error page instead
-	router.Use(handlers.RateLimitRequest)
+	router.Use(middleware.ProvideUserContext) // needed for everything else
+	router.Use(middleware.LogRequest)         // all pages need this
+	router.Use(middleware.SetPrivacyHeaders)  // all pages need this
+	router.Use(middleware.HandleError)        // if the inner handler fails, this shows the error page instead
+	router.Use(middleware.InitializeRateLimiter())
 
 	// watch and compile sass when in development mode
 	if config.GlobalConfig.InDevelopment {

server/handlers/limiter.go

@@ -1,98 +0,0 @@
-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/server/routes"
-)
-
-// CanRequestSkipLimiter determines if a request should bypass the rate limiter.
-// It exempts static assets and proxied image requests from rate limiting.
-func CanRequestSkipLimiter(r *http.Request) bool {
-	path := r.URL.Path
-	return strings.HasPrefix(path, "/img/") ||
-		strings.HasPrefix(path, "/css/") ||
-		strings.HasPrefix(path, "/js/") ||
-		strings.HasPrefix(path, "/proxy/s.pximg.net/")
-}
-
-// IPRateLimiter manages rate limiting on a per-IP basis.
-//
-// TODO: Should we put middlewares in a separate file?
-type IPRateLimiter struct {
-	ips     map[string]*rate.Limiter // Maps IP addresses to their respective rate limiters
-	mu      *sync.RWMutex            // Ensures thread-safe access to the map
-	limiter *rate.Limiter            // Global rate limiter used as a template for per-IP limiters
-}
-
-// NewIPRateLimiter creates a new instance of IPRateLimiter with the specified rate limit and burst.
-func NewIPRateLimiter(r rate.Limit, burst int) *IPRateLimiter {
-	return &IPRateLimiter{
-		ips:     make(map[string]*rate.Limiter),
-		mu:      &sync.RWMutex{},
-		limiter: rate.NewLimiter(r, burst),
-	}
-}
-
-// Allow checks if a request from the given IP is allowed based on the rate limit.
-// If the IP doesn't have a limiter, a new one is created.
-func (lim *IPRateLimiter) Allow(ip string) bool {
-	lim.mu.RLock()
-	rl, exists := lim.ips[ip]
-	lim.mu.RUnlock()
-
-	if !exists {
-		lim.mu.Lock()
-		rl, exists = lim.ips[ip]
-		if !exists {
-			// Create a new limiter for this IP using the global limiter's settings
-			rl = rate.NewLimiter(lim.limiter.Limit(), lim.limiter.Burst())
-			lim.ips[ip] = rl
-		}
-		lim.mu.Unlock()
-	}
-
-	return rl.Allow()
-}
-
-// Global rate limiter instance
-var limiter *IPRateLimiter
-
-// InitializeRateLimiter sets up the global rate limiter based on the application's configuration.
-// If the request limit is less than 1, it sets an infinite rate limit.
-func InitializeRateLimiter() {
-	r := float64(config.GlobalConfig.RequestLimit) / 30.0
-	if config.GlobalConfig.RequestLimit < 1 {
-		r = math.Inf(1)
-	}
-	limiter = NewIPRateLimiter(rate.Limit(r), 3)
-}
-
-// RateLimitRequest is a middleware that applies rate limiting to incoming HTTP requests.
-// It exempts certain requests (as defined by CanRequestSkipLimiter) from rate limiting.
-func RateLimitRequest(h http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		ip, _, _ := net.SplitHostPort(r.RemoteAddr)
-
-		if CanRequestSkipLimiter(r) {
-			h.ServeHTTP(w, r)
-			return
-		}
-
-		if !limiter.Allow(ip) {
-			// If the request exceeds the rate limit, return an HTTP 429 Too Many Requests error
-			routes.ErrorPage(w, r, errors.New("Too many requests"), http.StatusTooManyRequests)
-		} else {
-			// If the request is within the rate limit, proceed to the next handler
-			h.ServeHTTP(w, r)
-		}
-	})
-}

server/middleware/csp.go

@@ -1,4 +1,4 @@
-package handlers
+package middleware
 
 import (
 	"fmt"

server/middleware/doc.go

@@ -1,5 +1,5 @@
 /*
-Package handlers provides HTTP request handling functionality for PixivFE.
+Package middleware provides HTTP request handling functionality for PixivFE.
 It defines various middleware functions, request handlers, and routing logic to manage
 incoming HTTP requests and produce appropriate responses.
 
@@ -11,4 +11,4 @@ middlewares can be applied to routes to add cross-cutting functionality across m
 Route definitions are centralized in the DefineRoutes function, which sets up all paths
 and their corresponding handlers using the gorilla/mux router.
 */
-package handlers
+package middleware

server/middleware/error_handler.go

@@ -1,4 +1,4 @@
-package handlers
+package middleware
 
 import (
 	"bytes"

server/middleware/limiter.go

@@ -0,0 +1,75 @@
+package middleware
+
+import (
+	"log"
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/sethvargo/go-limiter/httplimit"
+	"github.com/sethvargo/go-limiter/memorystore"
+
+	"codeberg.org/vnpower/pixivfe/v2/config"
+)
+
+// CanRequestSkipLimiter determines if a request should bypass the rate limiter.
+// It exempts static assets and proxied image requests from rate limiting.
+func CanRequestSkipLimiter(r *http.Request) bool {
+	path := r.URL.Path
+	return strings.HasPrefix(path, "/img/") ||
+		strings.HasPrefix(path, "/css/") ||
+		strings.HasPrefix(path, "/js/") ||
+		strings.HasPrefix(path, "/proxy/s.pximg.net/")
+}
+
+// NewIPRateLimiter creates a new instance of IPRateLimiter with the specified rate limit and burst.
+//
+// ## Arguments
+//
+// Tokens: Number of tokens allowed per interval.
+// Interval: Interval until tokens reset.
+func NewIPRateLimiter(Tokens uint64, Interval time.Duration) (*httplimit.Middleware, error) {
+	store, err := memorystore.New(&memorystore.Config{
+		Tokens:   Tokens,
+		Interval: Interval,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return httplimit.NewMiddleware(store, httplimit.IPKeyFunc("X-Forwarded-For"))
+}
+
+// Global rate limiter instance
+var limiter *httplimit.Middleware
+
+// InitializeRateLimiter sets up the global rate limiter based on the application's configuration.
+// If the request limit is less than 1, it sets an infinite rate limit.
+// 
+// Returns the rate limit middleware
+func InitializeRateLimiter() func(http.Handler) http.Handler {
+	if config.GlobalConfig.RequestLimit < 1 {
+		limiter = nil
+	} else {
+		var err error
+		limiter, err = NewIPRateLimiter(config.GlobalConfig.RequestLimit, 30*time.Second)
+		if err != nil {
+			log.Panic(err)
+		}
+	}
+	return rateLimitRequest
+}
+
+// RateLimitRequest is a middleware that applies rate limiting to incoming HTTP requests.
+// It exempts certain requests (as defined by CanRequestSkipLimiter) from rate limiting.
+func rateLimitRequest(h http.Handler) http.Handler {
+	if limiter == nil {
+		return h
+	}
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if CanRequestSkipLimiter(r) {
+			h.ServeHTTP(w, r)
+			return
+		}
+		limiter.Handle(h).ServeHTTP(w, r)
+	})
+}

server/middleware/logger.go

@@ -1,4 +1,4 @@
-package handlers
+package middleware
 
 import (
 	"net/http"

server/middleware/recover.go

@@ -1,4 +1,4 @@
-package handlers
+package middleware
 
 import (
 	"errors"

server/middleware/router.go

@@ -1,4 +1,4 @@
-package handlers
+package middleware
 
 import (
 	"errors"
@@ -45,6 +45,7 @@ func DefineRoutes() *mux.Router {
 	}).HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		url := r.URL
 		url.Path = url.Path[0 : len(url.Path)-1]
+		// @iacore: i think this won't have open redirect vuln
 		http.Redirect(w, r, url.String(), http.StatusPermanentRedirect)
 	})
 

server/middleware/set_context.go

@@ -1,4 +1,4 @@
-package handlers
+package middleware
 
 import (
 	"net/http"