Merge pull request #1403 from bbx-winner/fix-upload-file

feat: add filename/type when uploading form-data

docs/CHANGELOG.md

@@ -1,5 +1,14 @@
 # Release History
 
+## v4.1.7 (2022-07-18)
+
+**go version**
+
+- fix: using '@FILEPATH' to indicate the path of the file 
+- feat: support indicating type and filename when uploading file
+- feat: support to infer MIME type of the file automatically
+
+
 ## v4.1.6 (2022-07-04)
 
 **go version**

hrp/internal/builtin/function.go

@@ -7,9 +7,13 @@ import (
 	"fmt"
 	"math"
 	"math/rand"
+	"mime"
 	"mime/multipart"
+	"net/textproto"
 	"os"
 	"path/filepath"
+	"regexp"
+	"strings"
 	"time"
 
 	"github.com/pkg/errors"
@@ -31,6 +35,15 @@ var Functions = map[string]interface{}{
 	"multipart_content_type": multipartContentType,
 }
 
+// upload file path must starts with @, like @\"PATH\" or @PATH
+var regexUploadFilePath = regexp.MustCompile(`^@(.*)`)
+
+var quoteEscaper = strings.NewReplacer("\\", "\\\\", `"`, "\\\"")
+
+func escapeQuotes(s string) string {
+	return quoteEscaper.Replace(s)
+}
+
 func init() {
 	rand.Seed(time.Now().UnixNano())
 }
@@ -60,57 +73,153 @@ func MD5(str string) string {
 	return hex.EncodeToString(hasher.Sum(nil))
 }
 
-type TFormWriter struct {
+type TFormDataWriter struct {
 	Writer  *multipart.Writer
 	Payload *bytes.Buffer
 }
 
-func multipartEncoder(formMap map[string]interface{}) *TFormWriter {
-	payload := &bytes.Buffer{}
-	writer := multipart.NewWriter(payload)
-	for formKey, formValue := range formMap {
-		formValueString := fmt.Sprintf("%v", formValue)
-		if err := writeFormDataFile(writer, formKey, formValueString); err == nil {
-			// form value is a file path
-			continue
-		}
-		// form value is not a file path, write as raw string
-		if err := writer.WriteField(formKey, formValueString); err != nil {
-			log.Info().Err(err).Msgf("failed to write field: %v=%v, ignore", formKey, formValue)
-		}
+func (w *TFormDataWriter) writeCustomText(formKey, formValue, formType, formFileName string) error {
+	if w.Writer == nil {
+		return errors.New("form-data writer not initialized")
 	}
-	if err := writer.Close(); err != nil {
+	h := make(textproto.MIMEHeader)
+	// text doesn't have Content-Type by default
+	if formType != "" {
+		h.Set("Content-Type", formType)
 	}
-	return &TFormWriter{
-		Writer:  writer,
-		Payload: payload,
+	// text doesn't have filename in Content-Disposition by default
+	if formFileName == "" {
+		h.Set("Content-Disposition",
+			fmt.Sprintf(`form-data; name="%s"`, escapeQuotes(formKey)))
+	} else {
+		h.Set("Content-Disposition",
+			fmt.Sprintf(`form-data; name="%s"; filename="%s"`,
+				escapeQuotes(formKey), escapeQuotes(formFileName)))
+	}
+	part, err := w.Writer.CreatePart(h)
+	if err != nil {
+		return err
 	}
-}
 
-func writeFormDataFile(writer *multipart.Writer, fName, fPath string) error {
-	var err error
-	fPath, err = filepath.Abs(fPath)
-	if err != nil {
-		log.Error().Err(err).Str("path", fPath).Msg("convert absolute path failed")
-		return err
-	}
-	if !IsFilePathExists(fPath) {
-		return errors.Errorf("file %v not existed", fPath)
-	}
-	file, err := os.ReadFile(fPath)
-	if err != nil {
-		log.Error().Err(err).Str("path", fPath).Msg("read file failed")
-		return err
-	}
-	formFile, err := writer.CreateFormFile(fName, filepath.Base(fPath))
-	if err != nil {
-		return err
-	}
-	_, err = formFile.Write(file)
+	_, err = part.Write([]byte(formValue))
 	return err
 }
 
-func multipartContentType(w *TFormWriter) string {
+func (w *TFormDataWriter) writeCustomFile(formKey, formValue, formType, formFileName string) error {
+	if w.Writer == nil {
+		return errors.New("form-data writer not initialized")
+	}
+	fPath, err := filepath.Abs(formValue)
+	if err != nil {
+		return err
+	}
+	file, err := os.ReadFile(fPath)
+	if err != nil {
+		return err
+	}
+
+	if formType == "" {
+		formType = inferFormType(formValue)
+	}
+	if formFileName == "" {
+		formFileName = filepath.Base(formValue)
+	}
+	h := make(textproto.MIMEHeader)
+	h.Set("Content-Type", formType)
+	h.Set("Content-Disposition",
+		fmt.Sprintf(`form-data; name="%s"; filename="%s"`,
+			escapeQuotes(formKey), escapeQuotes(formFileName)))
+	part, err := w.Writer.CreatePart(h)
+	if err != nil {
+		return err
+	}
+
+	_, err = part.Write(file)
+	return err
+}
+
+func inferFormType(formValue string) string {
+	extName := filepath.Ext(formValue)
+	formType := mime.TypeByExtension(extName)
+	if formType == "" {
+		// file without extension name
+		return "application/octet-stream"
+	}
+	if strings.HasPrefix(formType, "text") {
+		// text/... types have the charset parameter set to "utf-8" by default.
+		return strings.TrimSuffix(formType, "; charset=utf-8")
+	}
+	return formType
+}
+
+func multipartEncoder(formMap map[string]interface{}) (*TFormDataWriter, error) {
+	payload := &bytes.Buffer{}
+	writer := multipart.NewWriter(payload)
+	tFormWriter := &TFormDataWriter{
+		Writer:  writer,
+		Payload: payload,
+	}
+	// e.g. formMap: {"file": "@\"$upload_file\";type=text/foo"}
+	for formKey, formData := range formMap {
+		formDataString := fmt.Sprintf("%v", formData)
+		formItems := strings.Split(formDataString, ";")
+		var isFilePath bool
+		var formValue, formType, formFileName string
+		for _, formItem := range formItems {
+			if formItem == "" {
+				continue
+			}
+			equalSignIndex := strings.Index(formItem, "=")
+			// parse form value, e.g. @\"$upload_file\"
+			if equalSignIndex == -1 {
+				matchRes := regexUploadFilePath.FindStringSubmatch(formItem)
+				if len(matchRes) > 1 {
+					// formItem started with @, regarded as File path
+					isFilePath = true
+					formValue = strings.Trim(matchRes[1], "\"")
+				} else {
+					// formItem is not a valid File path, regarded as Text instead
+					formValue = strings.TrimSuffix(strings.TrimPrefix(formItem, "\""), "\"")
+				}
+				continue
+			}
+			// parse form option, e.g. type=text/plain
+			leftPart := strings.TrimSpace(formItem[:equalSignIndex])
+			var rightPart string
+			if equalSignIndex < len(formItem)-1 {
+				rightPart = strings.TrimSpace(formItem[equalSignIndex+1:])
+			}
+			if (strings.ToLower(leftPart) != "type" && strings.ToLower(leftPart) != "filename") || rightPart == "" {
+				formOption := fmt.Sprintf("%s=%s", leftPart, rightPart)
+				log.Warn().Msgf("invalid form option: %v, ignore", formOption)
+				continue
+			}
+			if strings.ToLower(leftPart) == "type" {
+				formType = rightPart
+			}
+			if strings.ToLower(leftPart) == "filename" {
+				formFileName = rightPart
+			}
+		}
+		if isFilePath {
+			if err := tFormWriter.writeCustomFile(formKey, formValue, formType, formFileName); err != nil {
+				log.Error().Err(err).Msgf("failed to write file: %v=@\"%v\", exit", formKey, formValue)
+				return nil, err
+			}
+			continue
+		}
+		if err := tFormWriter.writeCustomText(formKey, formValue, formType, formFileName); err != nil {
+			log.Error().Err(err).Msgf("failed to write text: %v=%v, ignore", formKey, formValue)
+			return nil, err
+		}
+	}
+	if err := writer.Close(); err != nil {
+		log.Error().Err(err).Msg("failed to close form-data writer")
+	}
+	return tFormWriter, nil
+}
+
+func multipartContentType(w *TFormDataWriter) string {
 	if w.Writer == nil {
 		return ""
 	}

hrp/response.go

@@ -130,7 +130,7 @@ func (v *responseObject) searchField(field string, variablesMapping map[string]i
 		var err error
 		result, err = v.parser.Parse(field, variablesMapping)
 		if err != nil {
-			log.Error().Str("filed name", field).Err(err).Msg("fail to parse field before search")
+			log.Error().Str("field name", field).Err(err).Msg("fail to parse field before search")
 		}
 	}
 	// search field using jmespath or regex if parsed field is still string and contains specified fieldTags

hrp/step_api.go

@@ -2,6 +2,7 @@ package hrp
 
 import (
 	"fmt"
+
 	"github.com/jinzhu/copier"
 	"github.com/rs/zerolog/log"
 
@@ -126,7 +127,7 @@ func extendWithAPI(testStep *TStep, overriddenStep *API) {
 	// merge & override request
 	testStep.Request = overriddenStep.Request
 	// init upload
-	if testStep.Request.Upload != nil {
+	if len(testStep.Request.Upload) != 0 {
 		initUpload(testStep)
 	}
 	// merge & override variables

hrp/step_request.go

@@ -247,7 +247,7 @@ func (r *requestBuilder) prepareBody(stepVariables map[string]interface{}) error
 		dataBytes = vv
 	case bytes.Buffer:
 		dataBytes = vv.Bytes()
-	case *builtin.TFormWriter:
+	case *builtin.TFormDataWriter:
 		dataBytes = vv.Payload.Bytes()
 	default: // unexpected body type
 		return errors.New("unexpected request body type")
@@ -268,7 +268,7 @@ func initUpload(step *TStep) {
 }
 
 func prepareUpload(parser *Parser, step *TStep, stepVariables map[string]interface{}) (err error) {
-	if step.Request.Upload == nil {
+	if len(step.Request.Upload) == 0 {
 		return
 	}
 	uploadMap, err := parser.Parse(step.Request.Upload, stepVariables)

hrp/testcase.go

@@ -148,7 +148,7 @@ func (path *TestCasePath) ToTestCase() (*TestCase, error) {
 			})
 		} else if step.Request != nil {
 			// init upload
-			if step.Request.Upload != nil {
+			if len(step.Request.Upload) != 0 {
 				initUpload(step)
 			}
 			testCase.TestSteps = append(testCase.TestSteps, &StepRequestWithOptionalArgs{

hrp/tests/upload_test.go

@@ -12,10 +12,10 @@ func TestCaseUploadFile(t *testing.T) {
 			SetBaseURL("https://httpbin.org").
 			WithVariables(map[string]interface{}{"upload_file": "test.env"}),
 		TestSteps: []hrp.IStep{
-			hrp.NewStep("upload file").
+			hrp.NewStep("upload file explicitly").
 				WithVariables(map[string]interface{}{
 					"m_encoder": "${multipart_encoder($m_upload)}",
-					"m_upload":  map[string]interface{}{"file": "$upload_file"},
+					"m_upload":  map[string]interface{}{"file": "@$upload_file"},
 				}).
 				POST("/post").
 				WithHeaders(map[string]string{"Content-Type": "${multipart_content_type($m_encoder)}"}).
@@ -23,12 +23,35 @@ func TestCaseUploadFile(t *testing.T) {
 				Validate().
 				AssertEqual("status_code", 200, "check status code").
 				AssertStartsWith("body.files.file", "UserName=test", "check uploaded file"),
-			hrp.NewStep("upload file with keyword").
+			hrp.NewStep("upload both text and file").
 				POST("/post").
-				WithUpload(map[string]interface{}{"file": "$upload_file"}).
+				WithUpload(map[string]interface{}{
+					"foo1":  "\"bar1\"",
+					"foo2":  "\"@$upload_file\"",
+					"foo3":  "\"\"@$upload_file\"\"",
+					"file1": "@\"$upload_file\"",
+					"file2": "@$upload_file",
+				}).
 				Validate().
 				AssertEqual("status_code", 200, "check status code").
-				AssertStartsWith("body.files.file", "UserName=test", "check uploaded file"),
+				AssertEqual("body.form.foo1", "bar1", "check foo1 in form").
+				AssertEqual("body.form.foo2", "@$upload_file", "check foo2 in form").
+				AssertEqual("body.form.foo3", "\"@$upload_file\"", "check foo3 in form").
+				AssertStartsWith("body.files.file1", "UserName=test", "check uploaded file1").
+				AssertStartsWith("body.files.file2", "UserName=test", "check uploaded file2"),
+			hrp.NewStep("upload empty field").
+				POST("/post").
+				WithUpload(map[string]interface{}{
+					"foo1":  "",
+					"foo2":  "\"\"",
+					"foo3":  "\"\";",
+					"dummy": ";filename=empty",
+				}).
+				Validate().
+				AssertEqual("status_code", 200, "check status code").
+				AssertEqual("body.form.foo1", "", "check foo1 in form").
+				AssertEqual("body.form.foo2", "", "check foo2 in form").
+				AssertEqual("body.files.dummy", "", "check dummy file in files"),
 		},
 	}