gophertype/pkg/ratelimit/ratelimit.go

140 lines
3.4 KiB
Go
Raw Normal View History

package ratelimit
import (
"fmt"
"strings"
"time"
"git.kirsle.net/apps/gophertype/pkg/cache"
"git.kirsle.net/apps/gophertype/pkg/constants"
)
// Limiter implements a Redis-backed rate limit for logins or otherwise.
type Limiter struct {
Namespace string // kind of rate limiter ("login")
ID interface{} // unique ID of the resource being pinged (str or ints)
Limit int // how many pings within the window period
Window time.Duration // the window period/expiration of Redis key
CooldownAt int // how many pings before the cooldown is enforced
Cooldown time.Duration // time to wait between fails
}
// Redis object behind the rate limiter.
type Data struct {
Pings int
NotBefore time.Time
}
// Ping the rate limiter.
func (l *Limiter) Ping() error {
var (
key = l.Key()
now = time.Now()
)
// Get stored data from Redis if any.
var data Data
cache.GetJSON(key, &data)
// Are we cooling down?
if now.Before(data.NotBefore) {
return fmt.Errorf(
"You are doing that too often. Please wait %s before trying again",
FormatDurationCoarse(data.NotBefore.Sub(now)),
)
}
// Increment the ping count.
data.Pings++
// Have we hit the wall?
if data.Pings >= l.Limit {
return fmt.Errorf(
"You have hit the rate limit; please wait the full %s before trying again",
FormatDurationCoarse(l.Window),
)
}
// Are we throttled?
if l.CooldownAt > 0 && data.Pings > l.CooldownAt {
data.NotBefore = now.Add(l.Cooldown)
if err := cache.SetJSON(key, data, l.Window); err != nil {
return fmt.Errorf("Couldn't set Redis key for rate limiter: %s", err)
}
return fmt.Errorf(
"Please wait %s before trying again. You have %d more attempt(s) remaining before you will be locked "+
"out for %s",
FormatDurationCoarse(l.Cooldown),
l.Limit-data.Pings,
FormatDurationCoarse(l.Window),
)
}
// Save their ping count to Redis.
if err := cache.SetJSON(key, data, l.Window); err != nil {
return fmt.Errorf("Couldn't set Redis key for rate limiter: %s", err)
}
return nil
}
// Clear the rate limiter, cleaning up the Redis key (e.g., after successful login).
func (l *Limiter) Clear() error {
return cache.Delete(l.Key())
}
// Key formats the Redis key.
func (l *Limiter) Key() string {
var str string
switch t := l.ID.(type) {
case int:
str = fmt.Sprintf("%d", t)
case uint64:
str = fmt.Sprintf("%d", t)
case int64:
str = fmt.Sprintf("%d", t)
case uint32:
str = fmt.Sprintf("%d", t)
case int32:
str = fmt.Sprintf("%d", t)
default:
str = fmt.Sprintf("%s", t)
}
return fmt.Sprintf(constants.RateLimitRedisKey, l.Namespace, str)
}
// FormatDurationCoarse returns a pretty printed duration with coarse granularity.
func FormatDurationCoarse(duration time.Duration) string {
var result = func(text string, v int64) string {
if v == 1 {
text = strings.TrimSuffix(text, "s")
}
return fmt.Sprintf(text, v)
}
if duration.Seconds() < 60.0 {
return result("%d seconds", int64(duration.Seconds()))
}
if duration.Minutes() < 60.0 {
return result("%d minutes", int64(duration.Minutes()))
}
if duration.Hours() < 24.0 {
return result("%d hours", int64(duration.Hours()))
}
days := int64(duration.Hours() / 24)
if days < 30 {
return result("%d days", days)
}
months := int64(days / 30)
if months < 12 {
return result("%d months", months)
}
years := int64(days / 365)
return result("%d years", years)
}