102 lines
2.3 KiB
Go
102 lines
2.3 KiB
Go
package server
|
|
|
|
import (
|
|
"net"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"uncensored-send/internal/config"
|
|
)
|
|
|
|
// limiter is a per-key token bucket. It is deliberately small: there is no
|
|
// dependency to pull in for this, and uploads are the only thing worth limiting.
|
|
type limiter struct {
|
|
rate float64 // tokens per second
|
|
burst float64
|
|
|
|
mu sync.Mutex
|
|
buckets map[string]*bucket
|
|
lastGC time.Time
|
|
}
|
|
|
|
type bucket struct {
|
|
tokens float64
|
|
last time.Time
|
|
}
|
|
|
|
func newLimiter(perHour, burst int) *limiter {
|
|
if burst < 1 {
|
|
burst = 1
|
|
}
|
|
return &limiter{
|
|
rate: float64(perHour) / 3600,
|
|
burst: float64(burst),
|
|
buckets: map[string]*bucket{},
|
|
}
|
|
}
|
|
|
|
// allow consumes a token for key, reporting whether one was available.
|
|
func (l *limiter) allow(key string, now time.Time) bool {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
|
|
b, ok := l.buckets[key]
|
|
if !ok {
|
|
b = &bucket{tokens: l.burst, last: now}
|
|
l.buckets[key] = b
|
|
}
|
|
b.tokens = min(l.burst, b.tokens+now.Sub(b.last).Seconds()*l.rate)
|
|
b.last = now
|
|
|
|
l.gcLocked(now)
|
|
|
|
if b.tokens < 1 {
|
|
return false
|
|
}
|
|
b.tokens--
|
|
return true
|
|
}
|
|
|
|
// gcLocked drops buckets that have refilled completely and so carry no state.
|
|
func (l *limiter) gcLocked(now time.Time) {
|
|
if now.Sub(l.lastGC) < 10*time.Minute {
|
|
return
|
|
}
|
|
l.lastGC = now
|
|
for k, b := range l.buckets {
|
|
if b.tokens >= l.burst && now.Sub(b.last) > time.Hour {
|
|
delete(l.buckets, k)
|
|
}
|
|
}
|
|
}
|
|
|
|
// clientIP identifies the caller for rate-limiting and logging.
|
|
//
|
|
// X-Forwarded-For is trusted only when the immediate peer is a configured
|
|
// proxy, and then only to skip past further trusted hops: the first
|
|
// untrusted address from the right is the real client. Without this the header
|
|
// is pure client input and rate limiting would be trivially bypassed.
|
|
func clientIP(r *http.Request, c *config.Config) string {
|
|
host, _, err := net.SplitHostPort(r.RemoteAddr)
|
|
if err != nil {
|
|
host = r.RemoteAddr
|
|
}
|
|
peer := net.ParseIP(host)
|
|
if peer == nil || !c.TrustsProxy(peer) {
|
|
return host
|
|
}
|
|
hops := strings.Split(r.Header.Get("X-Forwarded-For"), ",")
|
|
for i := len(hops) - 1; i >= 0; i-- {
|
|
ip := net.ParseIP(strings.TrimSpace(hops[i]))
|
|
if ip == nil {
|
|
break // malformed: stop trusting the rest of the chain
|
|
}
|
|
if !c.TrustsProxy(ip) {
|
|
return ip.String()
|
|
}
|
|
}
|
|
return host
|
|
}
|