Skip to content

Refactor createMultipartRequestBody to support custom content types a… #204

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jun 5, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
117 changes: 85 additions & 32 deletions httpclient/multipartrequest.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
"math"
"mime/multipart"
"net/http"
"net/textproto"
"os"
"path/filepath"
"strings"
Expand All @@ -22,7 +23,7 @@
)

// DoMultiPartRequest creates and executes a multipart/form-data HTTP request for file uploads and form fields.
func (c *Client) DoMultiPartRequest(method, endpoint string, files map[string]string, params map[string]string, out interface{}) (*http.Response, error) {
func (c *Client) DoMultiPartRequest(method, endpoint string, files map[string]string, params map[string]string, contentTypes map[string]string, headersMap map[string]http.Header, out interface{}) (*http.Response, error) {
log := c.Logger
ctx, cancel := context.WithCancel(context.Background())
defer cancel() // Ensure the context is canceled when the function returns
Expand Down Expand Up @@ -56,7 +57,7 @@

log.Debug("Executing multipart request", zap.String("method", method), zap.String("endpoint", endpoint))

body, contentType, err := createMultipartRequestBody(files, params, log)
body, contentType, err := createMultipartRequestBody(files, params, contentTypes, headersMap, log)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -105,54 +106,106 @@
return resp, response.HandleAPIErrorResponse(resp, log)
}

// createMultipartRequestBody creates a multipart request body with the provided files and form fields.
func createMultipartRequestBody(files map[string]string, params map[string]string, log logger.Logger) (*bytes.Buffer, string, error) {
// createMultipartRequestBody creates a multipart request body with the provided files and form fields, supporting custom content types and headers.
func createMultipartRequestBody(files map[string]string, params map[string]string, contentTypes map[string]string, headersMap map[string]http.Header, log logger.Logger) (*bytes.Buffer, string, error) {
body := &bytes.Buffer{}
writer := multipart.NewWriter(body)

for fieldName, filePath := range files {
file, err := os.Open(filePath)
if err != nil {
log.Error("Failed to open file", zap.String("filePath", filePath), zap.Error(err))
if err := addFilePart(writer, fieldName, filePath, contentTypes, headersMap, log); err != nil {
return nil, "", err
}
defer file.Close()
}

// Use only the filename in the Content-Disposition header
part, err := writer.CreateFormFile(fieldName, filepath.Base(filePath))
if err != nil {
log.Error("Failed to create form file", zap.String("fieldName", fieldName), zap.Error(err))
for key, val := range params {
if err := addFormField(writer, key, val, log); err != nil {
return nil, "", err
}
}

fileSize, err := file.Stat()
if err != nil {
log.Error("Failed to get file info", zap.String("filePath", filePath), zap.Error(err))
return nil, "", err
}
if err := writer.Close(); err != nil {
log.Error("Failed to close writer", zap.Error(err))

Check warning

Code scanning / gosec

Errors unhandled. Warning

Errors unhandled.
return nil, "", err
}

// Start logging the progress
progressLogger := logUploadProgress(fileSize.Size(), log)
return body, writer.FormDataContentType(), nil
}

// Chunk the file upload and log the progress
err = chunkFileUpload(file, part, log, progressLogger)
if err != nil {
log.Error("Failed to copy file content", zap.String("filePath", filePath), zap.Error(err))
return nil, "", err
}
// addFilePart adds a file part to the multipart writer with the provided field name and file path.
func addFilePart(writer *multipart.Writer, fieldName, filePath string, contentTypes map[string]string, headersMap map[string]http.Header, log logger.Logger) error {
file, err := os.Open(filePath)

Check failure

Code scanning / gosec

Potential file inclusion via variable Error

Potential file inclusion via variable
if err != nil {
log.Error("Failed to open file", zap.String("filePath", filePath), zap.Error(err))

Check warning

Code scanning / gosec

Errors unhandled. Warning

Errors unhandled.
return err
}
defer file.Close()

for key, val := range params {
_ = writer.WriteField(key, val)
contentType := "application/octet-stream"
if ct, ok := contentTypes[fieldName]; ok {
contentType = ct
}

err := writer.Close()
var partHeaders textproto.MIMEHeader
if h, ok := headersMap[fieldName]; ok {
partHeaders = CustomFormDataHeader(fieldName, filepath.Base(filePath), contentType, h)
} else {
partHeaders = FormDataHeader(fieldName, contentType)
}

part, err := writer.CreatePart(partHeaders)
if err != nil {
log.Error("Failed to close writer", zap.Error(err))
return nil, "", err
log.Error("Failed to create form file part", zap.String("fieldName", fieldName), zap.Error(err))

Check warning

Code scanning / gosec

Errors unhandled. Warning

Errors unhandled.
return err
}

return body, writer.FormDataContentType(), nil
fileSize, err := file.Stat()
if err != nil {
log.Error("Failed to get file info", zap.String("filePath", filePath), zap.Error(err))

Check warning

Code scanning / gosec

Errors unhandled. Warning

Errors unhandled.
return err
}

progressLogger := logUploadProgress(fileSize.Size(), log)
if err := chunkFileUpload(file, part, log, progressLogger); err != nil {
log.Error("Failed to copy file content", zap.String("filePath", filePath), zap.Error(err))

Check warning

Code scanning / gosec

Errors unhandled. Warning

Errors unhandled.
return err
}

return nil
}

// addFormField adds a form field to the multipart writer with the provided key and value.
func addFormField(writer *multipart.Writer, key, val string, log logger.Logger) error {
fieldWriter, err := writer.CreatePart(FormDataHeader(key, "text/plain"))
if err != nil {
log.Error("Failed to create form field", zap.String("key", key), zap.Error(err))

Check warning

Code scanning / gosec

Errors unhandled. Warning

Errors unhandled.
return err
}
if _, err := fieldWriter.Write([]byte(val)); err != nil {
log.Error("Failed to write form field", zap.String("key", key), zap.Error(err))

Check warning

Code scanning / gosec

Errors unhandled. Warning

Errors unhandled.
return err
}
return nil
}

// FormDataHeader creates a textproto.MIMEHeader for a form data field with the provided field name and content type.
func FormDataHeader(fieldname, contentType string) textproto.MIMEHeader {
header := textproto.MIMEHeader{}
header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"`, fieldname))
header.Set("Content-Type", contentType)
return header
}

// CustomFormDataHeader creates a textproto.MIMEHeader for a form data field with the provided field name, file name, content type, and custom headers.
func CustomFormDataHeader(fieldname, filename, contentType string, customHeaders http.Header) textproto.MIMEHeader {
header := textproto.MIMEHeader{}
header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"; filename="%s"`, fieldname, filename))
header.Set("Content-Type", contentType)
for key, values := range customHeaders {
for _, value := range values {
header.Add(key, value)
}
}
return header
}

// chunkFileUpload reads the file in chunks and writes it to the writer.
Expand Down Expand Up @@ -195,7 +248,7 @@
return nil
}

// trackUploadProgress logs the upload progress based on the percentage of the total upload.
// logUploadProgress logs the upload progress based on the percentage of the total upload.
func logUploadProgress(totalSize int64, log logger.Logger) func(int64) {
var uploadedSize int64
var lastLoggedPercentage float64
Expand Down
Loading