gosora/common/ratelimit.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)
}