This repository has been archived by the owner on May 20, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 2
/
http.go
89 lines (76 loc) · 2.06 KB
/
http.go
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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
package internal
import (
"fmt"
"io"
"math/rand"
"net/http"
"time"
)
func init() {
// It's such a low impact use case, there's no need for crypto/rand.
// Famous last words...
rand.Seed(time.Now().Unix())
}
const millis int = 1000
const minRate int = 500
const defaultJitter int = 2
// RateLimit accepts a jitter value (in seconds), uses it to create a new random
// timeout (in milliseconds) and adds a minimum rate (500ms).
// If jitter is less than one second, a default value will be used (2s).
func rateLimit(jitter int) time.Duration {
if jitter < 1 {
jitter = defaultJitter
}
i := rand.Intn(jitter*millis) + minRate
return time.Duration(i) * time.Millisecond
}
////////////////////////////////////////////////////////////////////////////////
type transport struct {
http.RoundTripper
}
func newTransport(dir string) *transport {
d := http.DefaultTransport
// The file protocol enables easier testing
d.(*http.Transport).RegisterProtocol("file", http.NewFileTransport(http.Dir(dir)))
return &transport{d}
}
func (t *transport) RoundTrip(r *http.Request) (*http.Response, error) {
r.Header.Set("User-Agent", "linux:"+GeneratorName+":"+GeneratorVersion+" ("+GeneratorSource+")")
return t.RoundTripper.RoundTrip(r)
}
type clientConf struct {
Timeout int
Jitter int
}
type client struct {
http *http.Client
conf clientConf
}
func newClient(conf clientConf) *client {
h := &http.Client{
Transport: newTransport("."),
Timeout: time.Duration(conf.Timeout) * time.Second,
}
return &client{
http: h,
conf: conf,
}
}
func (c *client) Get(path string) (io.ReadCloser, error) {
r, err := c.http.Get(path)
if err != nil {
return nil, err
}
if r.StatusCode < 200 || r.StatusCode > 299 {
// NOTE: program run shouldn't be long lived (like a server) so should
// be safe to ignore any .Close() errors
r.Body.Close() //#nosec G104
return nil, fmt.Errorf("bad response status: %s", r.Status)
}
return r.Body, nil
}
func (c *client) RateLimitedGet(path string) (io.ReadCloser, error) {
d := rateLimit(c.conf.Jitter)
time.Sleep(d)
return c.Get(path)
}