101 lines
2.2 KiB
Go
101 lines
2.2 KiB
Go
package common
|
|
|
|
import (
|
|
"errors"
|
|
"strconv"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
var ErrBadRateLimiter = errors.New("That rate limiter doesn't exist")
|
|
var ErrExceededRateLimit = errors.New("You're exceeding a rate limit. Please wait a while before trying again.")
|
|
|
|
// TODO: Persist rate limits to disk
|
|
type RateLimiter interface {
|
|
LimitIP(limit, ip string) error
|
|
LimitUser(limit string, user int) error
|
|
}
|
|
|
|
type RateData struct {
|
|
value int
|
|
floorTime int
|
|
}
|
|
|
|
type RateFence struct {
|
|
duration int
|
|
max int
|
|
}
|
|
|
|
// TODO: Optimise this by using something other than a string when possible
|
|
type RateLimit struct {
|
|
data map[string][]RateData
|
|
fences []RateFence
|
|
|
|
sync.RWMutex
|
|
}
|
|
|
|
func NewRateLimit(fences []RateFence) *RateLimit {
|
|
for i, fence := range fences {
|
|
fences[i].duration = fence.duration * 1000 * 1000 * 1000
|
|
}
|
|
return &RateLimit{data: make(map[string][]RateData), fences: fences}
|
|
}
|
|
|
|
func (l *RateLimit) Limit(name string, ltype int) error {
|
|
l.Lock()
|
|
defer l.Unlock()
|
|
|
|
data, ok := l.data[name]
|
|
if !ok {
|
|
data = make([]RateData, len(l.fences))
|
|
for i, _ := range data {
|
|
data[i] = RateData{0, int(time.Now().Unix())}
|
|
}
|
|
}
|
|
|
|
for i, field := range data {
|
|
fence := l.fences[i]
|
|
diff := int(time.Now().Unix()) - field.floorTime
|
|
|
|
if diff >= fence.duration {
|
|
field = RateData{0, int(time.Now().Unix())}
|
|
data[i] = field
|
|
}
|
|
|
|
if field.value > fence.max {
|
|
return ErrExceededRateLimit
|
|
}
|
|
|
|
field.value++
|
|
data[i] = field
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
type DefaultRateLimiter struct {
|
|
limits map[string]*RateLimit
|
|
}
|
|
|
|
func NewDefaultRateLimiter() *DefaultRateLimiter {
|
|
return &DefaultRateLimiter{map[string]*RateLimit{
|
|
"register": NewRateLimit([]RateFence{{int(time.Hour / 2), 1}}),
|
|
}}
|
|
}
|
|
|
|
func (l *DefaultRateLimiter) LimitIP(limit, ip string) error {
|
|
limiter, ok := l.limits[limit]
|
|
if !ok {
|
|
return ErrBadRateLimiter
|
|
}
|
|
return limiter.Limit(ip, 0)
|
|
}
|
|
|
|
func (l *DefaultRateLimiter) LimitUser(limit string, user int) error {
|
|
limiter, ok := l.limits[limit]
|
|
if !ok {
|
|
return ErrBadRateLimiter
|
|
}
|
|
return limiter.Limit(strconv.Itoa(user), 1)
|
|
}
|