1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
|
package http
import (
"github.com/tsenart/tb"
"net"
"net/http"
"time"
)
var byteThrottler = tb.NewThrottler(25 * time.Millisecond)
// ByteThrottledHandler wraps an http.Handler with per host byte throttling to
// the specified byte rate, responding with 429 when throttled.
func ByteThrottledHandler(h http.Handler, rate int64) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
host, _, _ := net.SplitHostPort(r.RemoteAddr)
if byteThrottler.Halt(host, r.ContentLength, rate) {
http.Error(w, "Too many requests", 429)
return
}
h.ServeHTTP(w, r)
})
}
var reqThrottler = tb.NewThrottler(5 * time.Millisecond)
// ReqThrottledHandler wraps an http.Handler with per host request throttling
// to the specified request rate, responding with 429 when throttled.
func ReqThrottledHandler(h http.Handler, rate int64) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
host, _, _ := net.SplitHostPort(r.RemoteAddr)
if reqThrottler.Halt(host, 1, rate) {
http.Error(w, "Too many requests", 429)
return
}
h.ServeHTTP(w, r)
})
}
|
