126 lines
3.5 KiB
Go
126 lines
3.5 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)
|
|
|
|
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)
|
|
}
|
|
}
|