job-search/internal/mock/findatender/ratelimit.go

package mock

import (
	"fmt"
	"net"
	"net/http"
	"strconv"
	"sync"
	"time"
)

// ---------- rate limiting ----------

// clientWindow is one client's rate limiting state.
type clientWindow struct {
	hour        time.Time // clock hour the anchor belongs to
	anchor      time.Time // first request of that hour
	windowStart time.Time // start of the current window
	count       int       // accepted requests in the current window
}

// RateLimiter limits each client (by remote host) to a number of requests per
// fixed window, answering the rest with 429 Too Many Requests and a
// Retry-After header, as the Find a Tender API does. A client's windows are
// counted from its first request in each clock hour.
type RateLimiter struct {
	mu      sync.Mutex
	limit   int
	window  time.Duration
	clients map[string]*clientWindow
	now     func() time.Time
}

// NewRateLimiter returns a RateLimiter allowing limit requests per window.
func NewRateLimiter(limit int, window time.Duration) *RateLimiter {
	return &RateLimiter{
		limit:   limit,
		window:  window,
		clients: map[string]*clientWindow{},
		now:     time.Now,
	}
}

// allow records a request for key and reports whether it may proceed. When it
// may not, it also returns the whole seconds until the client's next window.
func (rl *RateLimiter) allow(key string) (bool, int) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	now := rl.now().UTC()
	hour := now.Truncate(time.Hour)

	client, ok := rl.clients[key]
	if !ok || !client.hour.Equal(hour) {
		// First request of this clock hour: it anchors the windows.
		if !ok && len(rl.clients) > 10000 {
			rl.prune(hour)
		}
		client = &clientWindow{
			hour:        hour,
			anchor:      now,
			windowStart: now,
		}
		rl.clients[key] = client
	}

	// Move to whichever fixed window, counted from the anchor, now falls in.
	index := now.Sub(client.anchor) / rl.window
	windowStart := client.anchor.Add(index * rl.window)
	if !windowStart.Equal(client.windowStart) {
		client.windowStart = windowStart
		client.count = 0
	}

	if client.count < rl.limit {
		client.count++
		return true, 0
	}

	// The next window opens at the end of this one, or at the top of the
	// next hour if that comes first (because a new hour re-anchors).
	next := client.windowStart.Add(rl.window)
	if nextHour := hour.Add(time.Hour); nextHour.Before(next) {
		next = nextHour
	}

	seconds := int((next.Sub(now) + time.Second - 1) / time.Second) // round up
	if seconds < 1 {
		seconds = 1
	}

	return false, seconds
}

// prune drops clients whose state belongs to an earlier hour.
func (rl *RateLimiter) prune(hour time.Time) {
	for key, client := range rl.clients {
		if client.hour.Before(hour) {
			delete(rl.clients, key)
		}
	}
}

// clientKey produces an identifier for the client based on remote address.
// Could be updated to support proxy and forwarded-for headers.
func clientKey(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}

// Middleware attaches to http handler. Checks to see if a client is allowed
// to make a request in their current window or returns a "TooManyRequests"
// status and the time the client should wait before making another request.
func (rl *RateLimiter) Middleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		ok, retry := rl.allow(clientKey(r))
		if !ok {
			w.Header().Set("Retry-After", strconv.Itoa(retry))
			w.Header().Set("Content-Type", "text/plain; charset=utf-8")
			w.WriteHeader(http.StatusTooManyRequests)
			fmt.Fprintf( //nolint:errcheck // the client may have gone; nothing to do
				w,
				"Rate limit of %d exceeded. Please retry after %d seconds.",
				rl.limit,
				retry,
			)
			return
		}
		next(w, r)
	}
}