internal: centralize size/length constants; better DoS protection
...also adds a bounded custom event registry for cardinality control but I ran out of space in the commit message. Praise be to the long descriptions... Signed-off-by: NotAShelf <raf@notashelf.dev> Change-Id: Ic205f69804c7fb24c39fa84abdd9546b6a6a6964
This commit is contained in:
parent
371c5f3506
commit
c3b77696aa
3 changed files with 118 additions and 0 deletions
56
internal/aggregate/custom_events.go
Normal file
56
internal/aggregate/custom_events.go
Normal file
|
|
@ -0,0 +1,56 @@
|
||||||
|
package aggregate
|
||||||
|
|
||||||
|
import "sync"
|
||||||
|
|
||||||
|
// Maintains a bounded set of allowed custom event names
|
||||||
|
type CustomEventRegistry struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
events map[string]struct{}
|
||||||
|
maxEvents int
|
||||||
|
overflowCount int
|
||||||
|
}
|
||||||
|
|
||||||
|
// Creates a registry with a maximum number of unique event names
|
||||||
|
func NewCustomEventRegistry(maxEvents int) *CustomEventRegistry {
|
||||||
|
return &CustomEventRegistry{
|
||||||
|
events: make(map[string]struct{}, maxEvents),
|
||||||
|
maxEvents: maxEvents,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add attempts to add an event name to the registry
|
||||||
|
// Returns the event name if accepted, "other" if rejected due to limit
|
||||||
|
func (r *CustomEventRegistry) Add(eventName string) string {
|
||||||
|
// Fast path: check with read lock first
|
||||||
|
r.mu.RLock()
|
||||||
|
if _, exists := r.events[eventName]; exists {
|
||||||
|
r.mu.RUnlock()
|
||||||
|
return eventName
|
||||||
|
}
|
||||||
|
r.mu.RUnlock()
|
||||||
|
|
||||||
|
// Slow path: acquire write lock to add
|
||||||
|
r.mu.Lock()
|
||||||
|
defer r.mu.Unlock()
|
||||||
|
|
||||||
|
// Double-check after acquiring write lock
|
||||||
|
if _, exists := r.events[eventName]; exists {
|
||||||
|
return eventName
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check limit
|
||||||
|
if len(r.events) >= r.maxEvents {
|
||||||
|
r.overflowCount++
|
||||||
|
return "other"
|
||||||
|
}
|
||||||
|
|
||||||
|
r.events[eventName] = struct{}{}
|
||||||
|
return eventName
|
||||||
|
}
|
||||||
|
|
||||||
|
// Returns the number of events rejected due to limit
|
||||||
|
func (r *CustomEventRegistry) OverflowCount() int {
|
||||||
|
r.mu.RLock()
|
||||||
|
defer r.mu.RUnlock()
|
||||||
|
return r.overflowCount
|
||||||
|
}
|
||||||
9
internal/limits/constants.go
Normal file
9
internal/limits/constants.go
Normal file
|
|
@ -0,0 +1,9 @@
|
||||||
|
package limits
|
||||||
|
|
||||||
|
// Size limits for request processing
|
||||||
|
const (
|
||||||
|
MaxEventSize = 4 * 1024 // 4KB max event payload
|
||||||
|
MaxPathLen = 2048 // max path length
|
||||||
|
MaxRefLen = 2048 // max referrer length
|
||||||
|
MaxWidth = 10000 // max reasonable screen width
|
||||||
|
)
|
||||||
53
internal/ratelimit/limiter.go
Normal file
53
internal/ratelimit/limiter.go
Normal file
|
|
@ -0,0 +1,53 @@
|
||||||
|
package ratelimit
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Implements a simple token bucket rate limiter
|
||||||
|
type TokenBucket struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
tokens int
|
||||||
|
capacity int
|
||||||
|
refill int
|
||||||
|
interval time.Duration
|
||||||
|
lastFill time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
// Creates a rate limiter with specified capacity and refill rate capacity
|
||||||
|
func NewTokenBucket(capacity, refillPerInterval int, interval time.Duration) *TokenBucket {
|
||||||
|
return &TokenBucket{
|
||||||
|
tokens: capacity,
|
||||||
|
capacity: capacity,
|
||||||
|
refill: refillPerInterval,
|
||||||
|
interval: interval,
|
||||||
|
lastFill: time.Now(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Allow checks if a request should be allowed
|
||||||
|
func (tb *TokenBucket) Allow() bool {
|
||||||
|
tb.mu.Lock()
|
||||||
|
defer tb.mu.Unlock()
|
||||||
|
|
||||||
|
// Refill tokens based on elapsed time
|
||||||
|
now := time.Now()
|
||||||
|
elapsed := now.Sub(tb.lastFill)
|
||||||
|
if elapsed >= tb.interval {
|
||||||
|
periods := int(elapsed / tb.interval)
|
||||||
|
tb.tokens += periods * tb.refill
|
||||||
|
if tb.tokens > tb.capacity {
|
||||||
|
tb.tokens = tb.capacity
|
||||||
|
}
|
||||||
|
tb.lastFill = now.Add(-elapsed % tb.interval)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if we have tokens available
|
||||||
|
if tb.tokens > 0 {
|
||||||
|
tb.tokens--
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
return false
|
||||||
|
}
|
||||||
Loading…
Add table
Add a link
Reference in a new issue