optimize code

go.mod

@@ -3,9 +3,10 @@ module github.com/link1st/go-stress-testing
 go 1.16
 
 require (
-	github.com/golang/protobuf v1.5.3
+	github.com/mattn/go-shellwords v1.0.12
 	golang.org/x/net v0.23.0
 	golang.org/x/text v0.14.0
 	google.golang.org/grpc v1.56.3
+	google.golang.org/protobuf v1.30.0
 	layeh.com/radius v0.0.0-20210819152912-ad72663a72ab
 )

go.sum

@@ -795,6 +795,8 @@ github.com/lyft/protoc-gen-star v0.6.1/go.mod h1:TGAoBVkt8w7MPG72TrKIu85MIdXwDuz
 github.com/lyft/protoc-gen-star/v2 v2.0.1/go.mod h1:RcCdONR2ScXaYnQC5tUzxzlpA3WVYF7/opLeUgcQs/o=
 github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU=
 github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
+github.com/mattn/go-shellwords v1.0.12 h1:M2zGm7EW6UQJvDeQxo4T51eKPurbeFbe8WtebGE2xrk=
+github.com/mattn/go-shellwords v1.0.12/go.mod h1:EZzvwXDESEeg03EKmM+RmDnNOPKG4lLtQsUlTZDWQ8Y=
 github.com/mattn/go-sqlite3 v1.14.14/go.mod h1:NyWgC/yNuGj7Q9rpYnZvas74GogHl5/Z4A/KQRfk6bU=
 github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY=
 github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3/go.mod h1:RagcQ7I8IeTMnF8JTXieKnO4Z6JCsikNEzj0DwauVzE=

main.go

@@ -83,7 +83,8 @@ func main() {
 		return
 	}
 	debug := strings.ToLower(debugStr) == "true"
-	request, err := model.NewRequest(requestURL, verify, code, 0, debug, path, headers, body, maxCon, http2, keepalive, redirect)
+	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/curl_model.go

@@ -5,9 +5,12 @@ import (
 	"encoding/json"
 	"errors"
 	"io"
+	"net/http"
 	"os"
 	"strings"
 
+	"github.com/mattn/go-shellwords"
+
 	"github.com/link1st/go-stress-testing/helper"
 )
 
@@ -55,60 +58,50 @@ func ParseTheFile(path string) (curl *CURL, err error) {
 		err = errors.New("读取文件失败:" + err.Error())
 		return
 	}
-	data := string(dataBytes)
-	for len(data) > 0 {
-		if strings.HasPrefix(data, "curl") {
-			data = data[5:]
-		}
-		data = strings.TrimSpace(data)
-		var (
-			key   string
-			value string
-		)
-		index := strings.Index(data, " ")
-		if index <= 0 {
-			break
+	args, err := shellwords.Parse(string(dataBytes))
+	if err != nil {
+		err = errors.New("解析文件失败:" + err.Error())
+		return
+	}
+	args = argsTrim(args)
+	var key string
+	for _, arg := range args {
+		arg = removeSpaces(arg)
+		if arg == "" {
+			continue
 		}
-		key = strings.TrimSpace(data[:index])
-		data = data[index+1:]
-		data = strings.TrimSpace(data)
-		// url
-		if !strings.HasPrefix(key, "-") {
-			key = strings.Trim(key, "'")
-			curl.Data["curl"] = []string{key}
-			// 去除首尾空格
-			data = removeSpaces(data)
+		if isURL(arg) {
+			curl.Data[keyCurl] = append(curl.Data[keyCurl], arg)
+			key = ""
 			continue
 		}
-		if strings.HasPrefix(data, "-") {
+		if isKey(arg) {
+			key = arg
 			continue
 		}
-		var (
-			endSymbol = " "
-		)
-		if strings.HasPrefix(data, "'") {
-			endSymbol = "'"
-			data = data[1:]
+		curl.Data[key] = append(curl.Data[key], arg)
+	}
+	return
+}
+
+func argsTrim(args []string) []string {
+	result := make([]string, 0)
+	for _, arg := range args {
+		arg = strings.TrimSpace(arg)
+		if arg == "\n" {
+			continue
 		}
-		index = strings.Index(data, endSymbol)
-		if index <= -1 {
-			index = len(data)
-			// break
+		if strings.Contains(arg, "\n") {
+			arg = strings.ReplaceAll(arg, "\n", "")
 		}
-		value = data[:index]
-		if len(data) >= index+1 {
-			data = data[index+1:]
+		if strings.Index(arg, "-X") == 0 {
+			result = append(result, arg[0:2])
+			result = append(result, arg[2:])
 		} else {
-			data = ""
+			result = append(result, arg)
 		}
-		// 去除首尾空格
-		data = removeSpaces(data)
-		if key == "" {
-			continue
-		}
-		curl.Data[key] = append(curl.Data[key], value)
 	}
-	return
+	return result
 }
 
 func removeSpaces(data string) string {
@@ -121,15 +114,27 @@ func removeSpaces(data string) string {
 	return data
 }
 
+func isKey(data string) bool {
+	return strings.HasPrefix(data, "-") || strings.HasPrefix(data, keyCurl)
+}
+
+func isURL(data string) bool {
+	return strings.HasPrefix(data, "http://") || strings.HasPrefix(data, "https://")
+}
+
 // String string
 func (c *CURL) String() (url string) {
 	curlByte, _ := json.Marshal(c)
 	return string(curlByte)
 }
 
+const (
+	keyCurl = "curl"
+)
+
 // GetURL 获取url
 func (c *CURL) GetURL() (url string) {
-	keys := []string{"curl", "--url", "--location"}
+	keys := []string{keyCurl, "--url", "--location"}
 	value := c.getDataValue(keys)
 	if len(value) <= 0 {
 		return
@@ -146,25 +151,25 @@ func (c *CURL) GetMethod() (method string) {
 		return c.defaultMethod()
 	}
 	method = strings.ToUpper(value[0])
-	if helper.InArrayStr(method, []string{"GET", "POST", "PUT", "DELETE"}) {
+	if helper.InArrayStr(method, []string{http.MethodGet, http.MethodPost, http.MethodPut, http.MethodDelete}) {
 		return method
 	}
 	return c.defaultMethod()
 }
 
 // defaultMethod 获取默认方法
 func (c *CURL) defaultMethod() (method string) {
-	method = "GET"
+	method = http.MethodGet
 	body := c.GetBody()
 	if len(body) > 0 {
-		return "POST"
+		return http.MethodPost
 	}
 	return
 }
 
 // GetHeaders 获取请求头
 func (c *CURL) GetHeaders() (headers map[string]string) {
-	headers = make(map[string]string, 0)
+	headers = make(map[string]string)
 	keys := []string{"-H", "--header"}
 	value := c.getDataValue(keys)
 	for _, v := range value {