github.com/cloudwego/hertz@v0.9.3/pkg/protocol/multipart.go (about)

     1  /*
     2   * Copyright 2022 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   *
    16   * The MIT License (MIT)
    17   *
    18   * Copyright (c) 2015-present Aliaksandr Valialkin, VertaMedia, Kirill Danshin, Erik Dubbelboer, FastHTTP Authors
    19   *
    20   * Permission is hereby granted, free of charge, to any person obtaining a copy
    21   * of this software and associated documentation files (the "Software"), to deal
    22   * in the Software without restriction, including without limitation the rights
    23   * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    24   * copies of the Software, and to permit persons to whom the Software is
    25   * furnished to do so, subject to the following conditions:
    26   *
    27   * The above copyright notice and this permission notice shall be included in
    28   * all copies or substantial portions of the Software.
    29   *
    30   * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    31   * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    32   * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    33   * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    34   * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    35   * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    36   * THE SOFTWARE.
    37   *
    38   * This file may have been modified by CloudWeGo authors. All CloudWeGo
    39   * Modifications are Copyright 2022 CloudWeGo Authors.
    40   */
    41  
    42  package protocol
    43  
    44  import (
    45  	"fmt"
    46  	"io"
    47  	"mime/multipart"
    48  	"net/http"
    49  	"net/textproto"
    50  	"os"
    51  	"path/filepath"
    52  	"strings"
    53  
    54  	"github.com/cloudwego/hertz/pkg/common/bytebufferpool"
    55  	"github.com/cloudwego/hertz/pkg/common/utils"
    56  	"github.com/cloudwego/hertz/pkg/network"
    57  	"github.com/cloudwego/hertz/pkg/protocol/consts"
    58  )
    59  
    60  func ReadMultipartForm(r io.Reader, boundary string, size, maxInMemoryFileSize int) (*multipart.Form, error) {
    61  	// Do not care about memory allocations here, since they are tiny
    62  	// compared to multipart data (aka multi-MB files) usually sent
    63  	// in multipart/form-data requests.
    64  
    65  	if size <= 0 {
    66  		return nil, fmt.Errorf("form size must be greater than 0. Given %d", size)
    67  	}
    68  	lr := io.LimitReader(r, int64(size))
    69  	mr := multipart.NewReader(lr, boundary)
    70  	f, err := mr.ReadForm(int64(maxInMemoryFileSize))
    71  	if err != nil {
    72  		return nil, fmt.Errorf("cannot read multipart/form-data body: %s", err)
    73  	}
    74  	return f, nil
    75  }
    76  
    77  // WriteMultipartForm writes the given multipart form f with the given
    78  // boundary to w.
    79  func WriteMultipartForm(w io.Writer, f *multipart.Form, boundary string) error {
    80  	// Do not care about memory allocations here, since multipart
    81  	// form processing is slow.
    82  	if len(boundary) == 0 {
    83  		panic("BUG: form boundary cannot be empty")
    84  	}
    85  
    86  	mw := multipart.NewWriter(w)
    87  	if err := mw.SetBoundary(boundary); err != nil {
    88  		return fmt.Errorf("cannot use form boundary %q: %s", boundary, err)
    89  	}
    90  
    91  	// marshal values
    92  	for k, vv := range f.Value {
    93  		for _, v := range vv {
    94  			if err := mw.WriteField(k, v); err != nil {
    95  				return fmt.Errorf("cannot write form field %q value %q: %s", k, v, err)
    96  			}
    97  		}
    98  	}
    99  
   100  	// marshal files
   101  	for k, fvv := range f.File {
   102  		for _, fv := range fvv {
   103  			vw, err := mw.CreatePart(fv.Header)
   104  			zw := network.NewWriter(vw)
   105  			if err != nil {
   106  				return fmt.Errorf("cannot create form file %q (%q): %s", k, fv.Filename, err)
   107  			}
   108  			fh, err := fv.Open()
   109  			if err != nil {
   110  				return fmt.Errorf("cannot open form file %q (%q): %s", k, fv.Filename, err)
   111  			}
   112  			if _, err = utils.CopyZeroAlloc(zw, fh); err != nil {
   113  				return fmt.Errorf("error when copying form file %q (%q): %s", k, fv.Filename, err)
   114  			}
   115  			if err = fh.Close(); err != nil {
   116  				return fmt.Errorf("cannot close form file %q (%q): %s", k, fv.Filename, err)
   117  			}
   118  		}
   119  	}
   120  
   121  	if err := mw.Close(); err != nil {
   122  		return fmt.Errorf("error when closing multipart form writer: %s", err)
   123  	}
   124  
   125  	return nil
   126  }
   127  
   128  func MarshalMultipartForm(f *multipart.Form, boundary string) ([]byte, error) {
   129  	var buf bytebufferpool.ByteBuffer
   130  	if err := WriteMultipartForm(&buf, f, boundary); err != nil {
   131  		return nil, err
   132  	}
   133  	return buf.B, nil
   134  }
   135  
   136  func WriteMultipartFormFile(w *multipart.Writer, fieldName, fileName string, r io.Reader) error {
   137  	// Auto detect actual multipart content type
   138  	cbuf := make([]byte, 512)
   139  	size, err := r.Read(cbuf)
   140  	if err != nil && err != io.EOF {
   141  		return err
   142  	}
   143  
   144  	partWriter, err := w.CreatePart(CreateMultipartHeader(fieldName, fileName, http.DetectContentType(cbuf[:size])))
   145  	if err != nil {
   146  		return err
   147  	}
   148  
   149  	if _, err = partWriter.Write(cbuf[:size]); err != nil {
   150  		return err
   151  	}
   152  
   153  	_, err = io.Copy(partWriter, r)
   154  	return err
   155  }
   156  
   157  func CreateMultipartHeader(param, fileName, contentType string) textproto.MIMEHeader {
   158  	hdr := make(textproto.MIMEHeader)
   159  
   160  	var contentDispositionValue string
   161  	if len(strings.TrimSpace(fileName)) == 0 {
   162  		contentDispositionValue = fmt.Sprintf(`form-data; name="%s"`, param)
   163  	} else {
   164  		contentDispositionValue = fmt.Sprintf(`form-data; name="%s"; filename="%s"`,
   165  			param, fileName)
   166  	}
   167  	hdr.Set("Content-Disposition", contentDispositionValue)
   168  
   169  	if len(contentType) > 0 {
   170  		hdr.Set(consts.HeaderContentType, contentType)
   171  	}
   172  	return hdr
   173  }
   174  
   175  func AddFile(w *multipart.Writer, fieldName, path string) error {
   176  	file, err := os.Open(path)
   177  	if err != nil {
   178  		return err
   179  	}
   180  	defer file.Close()
   181  	return WriteMultipartFormFile(w, fieldName, filepath.Base(path), file)
   182  }
   183  
   184  func ParseMultipartForm(r io.Reader, request *Request, size, maxInMemoryFileSize int) error {
   185  	m, err := ReadMultipartForm(r, request.multipartFormBoundary, size, maxInMemoryFileSize)
   186  	if err != nil {
   187  		return err
   188  	}
   189  
   190  	request.multipartForm = m
   191  	return nil
   192  }
   193  
   194  func SetMultipartFormWithBoundary(req *Request, m *multipart.Form, boundary string) {
   195  	req.multipartForm = m
   196  	req.multipartFormBoundary = boundary
   197  }