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 }