This repository has been archived by the owner on May 21, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 14
/
attacker.go
201 lines (171 loc) · 4.89 KB
/
attacker.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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
package ultron
import (
"fmt"
"io"
"io/ioutil"
"net"
"net/http"
"time"
"github.com/valyala/fasthttp"
"go.uber.org/zap"
)
type (
// Attacker 定义一个事务、请求,需要确保该对象是Goroutine-safe
Attacker interface {
Name() string
Fire() error
}
// HTTPPrepareFunc 构造http.Request函数,需要调用方定义,由HTTPAttacker来发送
HTTPPrepareFunc func() (*http.Request, error)
// HTTPResponseCheck http.Response校验函数,可由调用方自定义,如果返回error,则视为请求失败
HTTPResponseCheck func(resp *http.Response, body []byte) error
// HTTPAttacker http协议的Attacker实现
HTTPAttacker struct {
Client *http.Client
Prepare HTTPPrepareFunc
name string
CheckChain []HTTPResponseCheck
}
// FastHTTPAttacker a http attacker base on fasthttp: https://github.com/valyala/fasthttp
FastHTTPAttacker struct {
Client *fasthttp.Client
Prepare FastHTTPPrepareFunc
name string
CheckChain []FastHTTPResponseCheck
}
// FastHTTPPrepareFunc 构造fasthttp.Request请求参数
FastHTTPPrepareFunc func(*fasthttp.Request) error
// FastHTTPResponseCheck check fasthttp.Response
FastHTTPResponseCheck func(*fasthttp.Response) error
)
var (
// DefaultHTTPClient 默认http.Client
// http://tleyden.github.io/blog/2016/11/21/tuning-the-go-http-client-library-for-load-testing/
DefaultHTTPClient = &http.Client{
Timeout: 90 * time.Second,
Transport: &http.Transport{
Proxy: nil,
DialContext: (&net.Dialer{
Timeout: 30 * time.Second,
KeepAlive: 30 * time.Second,
DualStack: true,
}).DialContext,
DisableKeepAlives: false,
MaxIdleConns: 2000,
MaxIdleConnsPerHost: 1000,
IdleConnTimeout: 30 * time.Second,
TLSHandshakeTimeout: 10 * time.Second,
ExpectContinueTimeout: 1 * time.Second,
},
}
// DefaultFastHTTPClient define the default fasthttp client use in FastHTTPAttacker
DefaultFastHTTPClient = &fasthttp.Client{
Name: "ultron",
MaxConnsPerHost: 2000,
MaxIdleConnDuration: 90 * time.Second,
ReadTimeout: 90 * time.Second,
WriteTimeout: 60 * time.Second,
}
)
// NewHTTPAttacker 创建一个新的HTTPAttacker对象
func NewHTTPAttacker(n string, p HTTPPrepareFunc, check ...HTTPResponseCheck) *HTTPAttacker {
return &HTTPAttacker{
Client: DefaultHTTPClient,
Prepare: p,
name: n,
CheckChain: check,
}
}
// Name 返回HTTPAttacker的名称
func (ha *HTTPAttacker) Name() string {
return ha.name
}
// Fire HTTPAttacker发起一次请求
func (ha *HTTPAttacker) Fire() error {
if ha.Prepare == nil {
panic("please implement Prepare() first")
}
req, err := ha.Prepare()
if err != nil {
Logger.Error("occur error on creating new http.Request object", zap.Error(err))
return err
}
resp, err := ha.Client.Do(req)
if err != nil {
Logger.Error("occur error on sending http request", zap.Error(err))
return err
}
if ha.CheckChain == nil || len(ha.CheckChain) == 0 {
io.Copy(ioutil.Discard, resp.Body) // no checker defined, discard body
resp.Body.Close()
return nil
}
body, err := ioutil.ReadAll(resp.Body)
if err != nil {
Logger.Error("occur error on receiving http response", zap.Error(err))
return err
}
resp.Body.Close()
for _, check := range ha.CheckChain {
if check == nil {
continue
}
if err = check(resp, body); err != nil {
return err
}
}
return nil
}
// CheckHTTPStatusCode 检查状态码是否>=400, 如果是则视为请求失败
func CheckHTTPStatusCode(resp *http.Response, body []byte) error {
if resp.StatusCode >= http.StatusBadRequest {
return fmt.Errorf("bad status code: %d", resp.StatusCode)
}
return nil
}
// CheckFastHTTPStatusCode check if status code >= 400
func CheckFastHTTPStatusCode(resp *fasthttp.Response) error {
if code := resp.StatusCode(); code >= http.StatusBadRequest {
return fmt.Errorf("bad status code: %d", code)
}
return nil
}
// NewFastHTTPAttacker return a new instance of FastHTTPAttacker
func NewFastHTTPAttacker(n string, p FastHTTPPrepareFunc, check ...FastHTTPResponseCheck) *FastHTTPAttacker {
return &FastHTTPAttacker{
Client: DefaultFastHTTPClient,
name: n,
Prepare: p,
CheckChain: check,
}
}
// Name return attacker's name
func (fa *FastHTTPAttacker) Name() string {
return fa.name
}
// Fire send request and check response
func (fa *FastHTTPAttacker) Fire() error {
if fa.Prepare == nil {
panic("please impl Prepare() first")
}
req := fasthttp.AcquireRequest()
resp := fasthttp.AcquireResponse()
defer fa.release(req, resp)
err := fa.Prepare(req)
if err != nil {
return err
}
if err := fa.Client.Do(req, resp); err != nil {
return err
}
for _, c := range fa.CheckChain {
if err := c(resp); err != nil {
return err
}
}
return nil
}
func (fa *FastHTTPAttacker) release(req *fasthttp.Request, resp *fasthttp.Response) {
fasthttp.ReleaseRequest(req)
fasthttp.ReleaseResponse(resp)
}