website/pkg/ratelimit/ratelimit.go

package ratelimit

import (
	"fmt"
	"time"

	"code.nonshy.com/nonshy/website/pkg/config"
	"code.nonshy.com/nonshy/website/pkg/redis"
	"code.nonshy.com/nonshy/website/pkg/utility"
)

// 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
	redis.Get(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.",
			utility.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.",
			utility.FormatDurationCoarse(l.Window),
		)
	}

	// Are we throttled?
	if l.CooldownAt > 0 && data.Pings > l.CooldownAt {
		data.NotBefore = now.Add(l.Cooldown)
		if err := redis.Set(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.",
			utility.FormatDurationCoarse(l.Cooldown),
			l.Limit-data.Pings,
			utility.FormatDurationCoarse(l.Window),
		)
	}

	// Save their ping count to Redis.
	if err := redis.Set(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 redis.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(config.RateLimitRedisKey, l.Namespace, str)
}