Add admin config page and runtime-tunable upload/rate-limit settings
This commit is contained in:
@@ -48,6 +48,8 @@ func (b *tokenBucket) allow() bool {
|
||||
|
||||
type ipClient struct {
|
||||
bucket *tokenBucket
|
||||
max int
|
||||
burst int
|
||||
lastSeen time.Time
|
||||
}
|
||||
|
||||
@@ -58,6 +60,12 @@ type ipClient struct {
|
||||
// burst: optional burst capacity (defaults to max if <=0)
|
||||
// ttl: how long to keep idle IP buckets around
|
||||
func RateLimitByIP(max int, per time.Duration, burst int, ttl time.Duration) gin.HandlerFunc {
|
||||
return RateLimitByIPDynamic(func() int { return max }, per, func() int { return burst }, ttl)
|
||||
}
|
||||
|
||||
// RateLimitByIPDynamic is like RateLimitByIP but reads max/burst dynamically.
|
||||
// This allows changing limits at runtime (e.g. from an admin config page).
|
||||
func RateLimitByIPDynamic(maxFn func() int, per time.Duration, burstFn func() int, ttl time.Duration) gin.HandlerFunc {
|
||||
var (
|
||||
mu sync.Mutex
|
||||
clients = make(map[string]*ipClient)
|
||||
@@ -86,16 +94,22 @@ func RateLimitByIP(max int, per time.Duration, burst int, ttl time.Duration) gin
|
||||
}
|
||||
}
|
||||
|
||||
getClient := func(ip string, now time.Time) *ipClient {
|
||||
getClient := func(ip string, now time.Time, max int, burst int) *ipClient {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
c, ok := clients[ip]
|
||||
if !ok {
|
||||
c = &ipClient{bucket: newTokenBucket(max, per, burst), lastSeen: now}
|
||||
c = &ipClient{bucket: newTokenBucket(max, per, burst), max: max, burst: burst, lastSeen: now}
|
||||
clients[ip] = c
|
||||
return c
|
||||
}
|
||||
c.lastSeen = now
|
||||
// If settings changed, reset the bucket.
|
||||
if c.max != max || c.burst != burst {
|
||||
c.bucket = newTokenBucket(max, per, burst)
|
||||
c.max = max
|
||||
c.burst = burst
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
@@ -104,7 +118,15 @@ func RateLimitByIP(max int, per time.Duration, burst int, ttl time.Duration) gin
|
||||
cleanup(now)
|
||||
|
||||
ip := c.ClientIP()
|
||||
client := getClient(ip, now)
|
||||
max := maxFn()
|
||||
if max <= 0 {
|
||||
max = 1
|
||||
}
|
||||
burst := burstFn()
|
||||
if burst <= 0 {
|
||||
burst = max
|
||||
}
|
||||
client := getClient(ip, now, max, burst)
|
||||
|
||||
if !client.bucket.allow() {
|
||||
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
|
||||
|
||||
Reference in New Issue
Block a user