Franz Kafka
5b942a5fd6
Trim or remove comments that: - State the obvious (function names already convey purpose) - Repeat what the code clearly shows - Are excessively long without adding value Keep comments that explain *why*, not *what*.
146 lines
3.1 KiB
Go
146 lines
3.1 KiB
Go
// kafka — a privacy-respecting metasearch engine
|
|
// Copyright (C) 2026-present metamorphosis-dev
|
|
//
|
|
// This program is free software: you can redistribute it and/or modify
|
|
// it under the terms of the GNU Affero General Public License as published by
|
|
// the Free Software Foundation, either version 3 of the License, or
|
|
// (at your option) any later version.
|
|
//
|
|
// This program is distributed in the hope that it will be useful,
|
|
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
// GNU Affero General Public License for more details.
|
|
//
|
|
// You should have received a copy of the GNU Affero General Public License
|
|
// along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
|
|
package middleware
|
|
|
|
import (
|
|
"net"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"log/slog"
|
|
)
|
|
|
|
type RateLimitConfig struct {
|
|
Requests int
|
|
Window time.Duration
|
|
CleanupInterval time.Duration
|
|
}
|
|
|
|
func RateLimit(cfg RateLimitConfig, logger *slog.Logger) func(http.Handler) http.Handler {
|
|
requests := cfg.Requests
|
|
if requests <= 0 {
|
|
requests = 30
|
|
}
|
|
|
|
window := cfg.Window
|
|
if window <= 0 {
|
|
window = time.Minute
|
|
}
|
|
|
|
cleanup := cfg.CleanupInterval
|
|
if cleanup <= 0 {
|
|
cleanup = 5 * time.Minute
|
|
}
|
|
|
|
if logger == nil {
|
|
logger = slog.Default()
|
|
}
|
|
|
|
limiter := &ipLimiter{
|
|
requests: requests,
|
|
window: window,
|
|
clients: make(map[string]*bucket),
|
|
logger: logger,
|
|
}
|
|
|
|
go limiter.cleanup(cleanup)
|
|
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
ip := extractIP(r)
|
|
|
|
if !limiter.allow(ip) {
|
|
retryAfter := int(limiter.window.Seconds())
|
|
w.Header().Set("Retry-After", strconv.Itoa(retryAfter))
|
|
w.Header().Set("Content-Type", "text/plain; charset=utf-8")
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
_, _ = w.Write([]byte("429 Too Many Requests\n"))
|
|
logger.Debug("rate limited", "ip", ip)
|
|
return
|
|
}
|
|
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
}
|
|
|
|
type bucket struct {
|
|
count int
|
|
expireAt time.Time
|
|
}
|
|
|
|
type ipLimiter struct {
|
|
requests int
|
|
window time.Duration
|
|
clients map[string]*bucket
|
|
mu sync.Mutex
|
|
logger *slog.Logger
|
|
}
|
|
|
|
func (l *ipLimiter) allow(ip string) bool {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
|
|
now := time.Now()
|
|
b, ok := l.clients[ip]
|
|
|
|
if !ok || now.After(b.expireAt) {
|
|
l.clients[ip] = &bucket{
|
|
count: 1,
|
|
expireAt: now.Add(l.window),
|
|
}
|
|
return true
|
|
}
|
|
|
|
b.count++
|
|
return b.count <= l.requests
|
|
}
|
|
|
|
func (l *ipLimiter) cleanup(interval time.Duration) {
|
|
ticker := time.NewTicker(interval)
|
|
defer ticker.Stop()
|
|
|
|
for range ticker.C {
|
|
l.mu.Lock()
|
|
now := time.Now()
|
|
for ip, b := range l.clients {
|
|
if now.After(b.expireAt) {
|
|
delete(l.clients, ip)
|
|
}
|
|
}
|
|
l.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
func extractIP(r *http.Request) string {
|
|
if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
|
|
parts := strings.SplitN(xff, ",", 2)
|
|
return strings.TrimSpace(parts[0])
|
|
}
|
|
if rip := r.Header.Get("X-Real-IP"); rip != "" {
|
|
return strings.TrimSpace(rip)
|
|
}
|
|
|
|
host, _, err := net.SplitHostPort(r.RemoteAddr)
|
|
if err != nil {
|
|
return r.RemoteAddr
|
|
}
|
|
return host
|
|
}
|