Merge pull request #149 from beihai0xff/master

新增参数 redirect

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

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
 }

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))

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
 }