diff --git a/main.go b/main.go index 2cd3bf4..cb08ee5 100644 --- a/main.go +++ b/main.go @@ -46,6 +46,7 @@ var ( keepalive = false // 是否开启长连接 cpuNumber = 1 // CUP 核数,默认为一核,一般场景下单核已经够用了 timeout int64 = 0 // 超时时间,默认不设置 + redirect = true // 是否重定向 ) func init() { @@ -63,6 +64,7 @@ func init() { flag.BoolVar(&keepalive, "k", keepalive, "是否开启长连接") flag.IntVar(&cpuNumber, "cpuNumber", cpuNumber, "CUP 核数,默认为一核") flag.Int64Var(&timeout, "timeout", timeout, "超时时间 单位 秒,默认不设置") + flag.BoolVar(&redirect, "redirect", redirect, "是否重定向") // 解析参数 flag.Parse() } @@ -81,7 +83,7 @@ func main() { return } debug := strings.ToLower(debugStr) == "true" - request, err := model.NewRequest(requestURL, verify, code, 0, debug, path, headers, body, maxCon, http2, keepalive) + request, err := model.NewRequest(requestURL, verify, code, 0, debug, path, headers, body, maxCon, http2, keepalive, redirect) if err != nil { fmt.Printf("参数不合法 %v \n", err) return diff --git a/model/request_model.go b/model/request_model.go index b464eee..9fa8329 100644 --- a/model/request_model.go +++ b/model/request_model.go @@ -86,6 +86,7 @@ type Request struct { HTTP2 bool // 是否使用http2.0 Keepalive bool // 是否开启长连接 Code int // 验证的状态码 + Redirect bool // 是否重定向 } // GetBody 获取请求数据 @@ -123,7 +124,7 @@ func (r *Request) GetVerifyWebSocket() VerifyWebSocket { // debug 是否开启debug // path curl文件路径 http接口压测,自定义参数设置 func NewRequest(url string, verify string, code int, timeout time.Duration, debug bool, path string, - reqHeaders []string, reqBody string, maxCon int, http2 bool, keepalive bool) (request *Request, err error) { + reqHeaders []string, reqBody string, maxCon int, http2, keepalive, redirect bool) (request *Request, err error) { var ( method = "GET" headers = make(map[string]string) @@ -200,6 +201,7 @@ func NewRequest(url string, verify string, code int, timeout time.Duration, debu HTTP2: http2, Keepalive: keepalive, Code: code, + Redirect: redirect, } return } diff --git a/server/client/http_client.go b/server/client/http_client.go index 16047bc..1cf2c7f 100644 --- a/server/client/http_client.go +++ b/server/client/http_client.go @@ -80,6 +80,12 @@ func HTTPRequest(chanID uint64, request *model.Request) (resp *http.Response, re Transport: tr, Timeout: timeout, } + if !request.Redirect { + client.CheckRedirect = func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + } + } + startTime := time.Now() resp, err = client.Do(req) requestTime = uint64(helper.DiffNano(startTime)) diff --git a/server/client/http_longclinet/long_clinet.go b/server/client/http_longclinet/long_clinet.go index 556ddc5..999c7f7 100644 --- a/server/client/http_longclinet/long_clinet.go +++ b/server/client/http_longclinet/long_clinet.go @@ -67,7 +67,16 @@ func createLangHTTPClient(request *model.Request) *http.Client { } _ = http2.ConfigureTransport(tr) } - return &http.Client{ + + client := &http.Client{ Transport: tr, } + + if !request.Redirect { + client.CheckRedirect = func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + } + } + + return client }