bwilliams
465358609a
Short and single-letter names in application code become full words (tests keep theirs). Wrapped calls, signatures and composite literals put one item per line, and one-line function bodies are expanded. Exported names and groups of constants get doc comments, and comments refer to the renamed identifiers.
130 lines
3.7 KiB
Go
130 lines
3.7 KiB
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
|
|
if windowStart := client.anchor.Add(index * rl.window); !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(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)
|
|
}
|
|
}
|