File/package naming convention conforming better to go standards - as ugly as it is

internal/mock/findatender/ratelimit.go

@@ -0,0 +1,125 @@
+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)
+
+	c, ok := rl.clients[key]
+	if !ok || !c.hour.Equal(hour) {
+		// First request of this clock hour: it anchors the windows.
+		if !ok && len(rl.clients) > 10000 {
+			rl.prune(hour)
+		}
+		c = &clientWindow{hour: hour, anchor: now, windowStart: now}
+		rl.clients[key] = c
+	}
+
+	// Move to whichever fixed window, counted from the anchor, now falls in.
+	idx := now.Sub(c.anchor) / rl.window
+	if ws := c.anchor.Add(idx * rl.window); !ws.Equal(c.windowStart) {
+		c.windowStart = ws
+		c.count = 0
+	}
+
+	if c.count < rl.limit {
+		c.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 := c.windowStart.Add(rl.window)
+	if nh := hour.Add(time.Hour); nh.Before(next) {
+		next = nh
+	}
+	secs := int((next.Sub(now) + time.Second - 1) / time.Second) // round up
+	if secs < 1 {
+		secs = 1
+	}
+	return false, secs
+}
+
+// prune drops clients whose state belongs to an earlier hour.
+func (rl *RateLimiter) prune(hour time.Time) {
+	for k, c := range rl.clients {
+		if c.hour.Before(hour) {
+			delete(rl.clients, k)
+		}
+	}
+}
+
+// clientKey produces an identifier for the client based on remote address.
+// Could be updated to support proxy and forwarded-for headers.
+func clientKey(req *http.Request) string {
+	host, _, err := net.SplitHostPort(req.RemoteAddr)
+	if err != nil {
+		return req.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, req *http.Request) {
+		ok, retry := rl.allow(clientKey(req))
+		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, req)
+	}
+}