grokify/mogo

// Package multipartutil provides helper functionality for using multipart.Writer.
// Steps are to call NewMultipartBuilder(), write parts, call builder.Close(), and
// retrieve Content-Type header from builder.Writer.FormDataContentType().
package multipartutil

import (
    "bytes"
    "encoding/base64"
    "encoding/json"
    "fmt"
    "io"
    "mime"
    "mime/multipart"
    "net/http"
    "net/textproto"
    "net/url"
    "os"
    "path/filepath"
    "strings"

    hum "github.com/grokify/mogo/net/http/httputilmore"
)

func NewReaderBodyBytes(body []byte, boundary string) *multipart.Reader {
    return multipart.NewReader(bytes.NewReader(body), boundary)
}

func NewMultipartReaderForHTTPResponse(resp *http.Response) (*multipart.Reader, error) {
    contentType := resp.Header.Get(hum.HeaderContentType)
    mediaType, params, err := mime.ParseMediaType(contentType)
    if err != nil {
        return nil, err
    } else if !strings.HasPrefix(mediaType, "multipart/") {
        return nil, fmt.Errorf("MediaType is not multipart [%v]", mediaType)
    }
    if boundary, ok := params["boundary"]; !ok {
        return nil, fmt.Errorf("MIME Boundary not found in Content-Type header [%v]", contentType)
    } else {
        return multipart.NewReader(resp.Body, boundary), nil
    }
}

// MultipartBuilder is a multipart helper.
type MultipartBuilder struct {
    Buffer *bytes.Buffer
    Writer *multipart.Writer
}

// NewMultipartBuilder instantiates a new MultipartBuilder.
func NewMultipartBuilder() MultipartBuilder {
    builder := MultipartBuilder{}
    var b bytes.Buffer
    builder.Buffer = &b
    builder.Writer = multipart.NewWriter(&b)
    return builder
}

// WriteURLValues writes simple header key value strings using `url.Values`
// as an input parameter.
func (builder *MultipartBuilder) WriteURLValues(params url.Values) error {
    for key, vals := range params {
        for _, val := range vals {
            err := builder.WriteFieldString(key, val)
            if err != nil {
                return err
            }
        }
    }
    return nil
}

// WriteFieldString adds a text part.
func (builder *MultipartBuilder) WriteFieldString(partName string, data string) error {
    return builder.Writer.WriteField(partName, data)
}

// WriteFieldAsJSON adds a JSON part.
func (builder *MultipartBuilder) WriteFieldAsJSON(partName string, data any, base64Encode bool) error {
    jsonBytes, err := json.Marshal(data)
    if err != nil {
        return err
    }

    header := textproto.MIMEHeader{}
    header.Add(hum.HeaderContentDisposition, fmt.Sprintf(`form-data; name="%s"`, partName))
    header.Add(hum.HeaderContentType, hum.ContentTypeAppJSONUtf8)
    if base64Encode {
        header.Add(hum.HeaderContentTransferEncoding, "base64")
    }

    partWriter, err := builder.Writer.CreatePart(header)
    if err != nil {
        return err
    }

    if base64Encode {
        str := base64.StdEncoding.EncodeToString(jsonBytes)
        _, err = bytes.NewBuffer([]byte(str)).WriteTo(partWriter)
    } else {
        _, err = bytes.NewBuffer(jsonBytes).WriteTo(partWriter)
    }
    return err
}

// WriteFilepathPlus adds a file part given a filename with the Content Type and
// other associated headers as needed. After builder.Close() has been called,
// use like `req, err := http.NewRequest("POST", url, builder.Buffer)`.
// Content-Disposition uses optional attribute as defined here:
// https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Disposition
func (builder *MultipartBuilder) WriteFilePathPlus(partName, srcFilepath string, base64Encode bool) error {
    partName = strings.TrimSpace(partName)
    _, filename := filepath.Split(srcFilepath)
    filename = strings.TrimSpace(filename)
    mimeType := mime.TypeByExtension(filepath.Ext(srcFilepath))

    header := textproto.MIMEHeader{}
    cd := []string{"form-data"}
    if len(partName) > 0 {
        cd = append(cd, fmt.Sprintf(`name="%s"`, partName))
    }
    if len(filename) > 0 {
        cd = append(cd, fmt.Sprintf(`filename="%s"`, filename))
    }
    header.Add(hum.HeaderContentDisposition, strings.Join(cd, "; "))

    if len(mimeType) > 0 {
        header.Add(hum.HeaderContentType, mimeType)
    }
    if base64Encode {
        header.Add(hum.HeaderContentTransferEncoding, "base64")
    }

    partWriter, err := builder.Writer.CreatePart(header)
    if err != nil {
        return err
    }

    fileBytes, err := os.ReadFile(srcFilepath)
    if err != nil {
        return err
    }

    if base64Encode {
        str := base64.StdEncoding.EncodeToString(fileBytes)
        _, err = bytes.NewBuffer([]byte(str)).WriteTo(partWriter)
    } else {
        _, err = bytes.NewBuffer(fileBytes).WriteTo(partWriter)
    }
    return err
}

// WriteFilePath adds a file part given a filename.
func (builder *MultipartBuilder) WriteFilePath(partName, srcFilepath string) error {
    file, err := os.Open(srcFilepath)
    if err != nil {
        return err
    }
    defer file.Close()
    _, filename := filepath.Split(srcFilepath)
    return builder.WriteFileReader(partName, filename, file)
}

// WriteFileHeader adds a file part given a part name and *multipart.FileHeader.
// See more at http://sanatgersappa.blogspot.com/2013/03/handling-multiple-file-uploads-in-go.html
// and https://gist.github.com/sanatgersappa/5127317#file-app-go
func (builder *MultipartBuilder) WriteFileHeader(partName string, fileHeader *multipart.FileHeader) error {
    file, err := fileHeader.Open()
    if err != nil {
        return err
    }
    defer file.Close()
    return builder.WriteFileReader(partName, fileHeader.Filename, file)
}

// WriteFileReader adds a file part given a filename and `io.Reader`.
func (builder *MultipartBuilder) WriteFileReader(partName, filename string, src io.Reader) error {
    fw, err := builder.Writer.CreateFormFile(partName, filename)
    if err != nil {
        return err
    }
    _, err = io.Copy(fw, src)
    return err
}

// Close closes the `multipart.Writer`.
func (builder *MultipartBuilder) Close() error {
    return builder.Writer.Close()
}

// ContentType returns the content type for the `Content-Type` header.
func (builder *MultipartBuilder) ContentType() string {
    return builder.Writer.FormDataContentType()
}

// String returns the MIME parts as a string.
func (builder *MultipartBuilder) String() string {
    return builder.Buffer.String()
}