josie / simplegit

package web

import (
	"sync"
	"time"
)

// ipLimiter is a small in-memory sliding-window rate limiter keyed by client
// IP. It guards unauthenticated guest writes; a periodic sweep inside allow
// drops keys whose hits have all aged out.
type ipLimiter struct {
	mu     sync.Mutex
	hits   map[string][]time.Time
	limit  int
	window time.Duration
	calls  int
}

func newIPLimiter(limit int, window time.Duration) *ipLimiter {
	return &ipLimiter{hits: make(map[string][]time.Time), limit: limit, window: window}
}

// allow records a hit for key and reports whether it is within the limit.
func (l *ipLimiter) allow(key string) bool {
	now := time.Now()
	cut := now.Add(-l.window)

	l.mu.Lock()
	defer l.mu.Unlock()

	l.calls++
	if l.calls%512 == 0 {
		l.sweep(cut)
	}

	recent := l.hits[key][:0]
	for _, t := range l.hits[key] {
		if t.After(cut) {
			recent = append(recent, t)
		}
	}
	if len(recent) >= l.limit {
		l.hits[key] = recent
		return false
	}
	l.hits[key] = append(recent, now)
	return true
}

// sweep drops keys whose recorded hits have all left the window. Called
// with the mutex held.
func (l *ipLimiter) sweep(cut time.Time) {
	for key, times := range l.hits {
		recent := times[:0]
		for _, t := range times {
			if t.After(cut) {
				recent = append(recent, t)
			}
		}
		if len(recent) == 0 {
			delete(l.hits, key)
			continue
		}
		l.hits[key] = recent
	}
}

// reset forgets key's recorded hits, so a limiter keyed on credentials
// rewards a success (e.g. a signed-in user clears the failed-login window).
func (l *ipLimiter) reset(key string) {
	l.mu.Lock()
	defer l.mu.Unlock()
	delete(l.hits, key)
}