-
Notifications
You must be signed in to change notification settings - Fork 53
/
multipart.go
139 lines (117 loc) · 3.22 KB
/
multipart.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
package multipart
import (
"bytes"
"errors"
"io"
"io/ioutil"
"mime/multipart"
"strconv"
"strings"
c "gopkg.in/h2non/gentleman.v1/context"
p "gopkg.in/h2non/gentleman.v1/plugin"
)
// Values represents multiple multipart from values.
type Values []string
// DataFields represents a map of text based fields.
type DataFields map[string]Values
// FormFile represents the file form field data.
type FormFile struct {
Name string
Reader io.Reader
}
// FormData represents the supported form fields by file and string data.
type FormData struct {
Data DataFields
Files []FormFile
}
// File creates a new multipart form based on a unique file field
// from the given io.ReadCloser stream.
func File(name string, reader io.Reader) p.Plugin {
return p.NewRequestPlugin(func(ctx *c.Context, h c.Handler) {
file := FormFile{name, reader}
data := FormData{Files: []FormFile{file}}
handle(ctx, h, data)
})
}
// Files creates a multipart form based on files fields.
func Files(files []FormFile) p.Plugin {
return p.NewRequestPlugin(func(ctx *c.Context, h c.Handler) {
data := FormData{Files: files}
handle(ctx, h, data)
})
}
// Fields creates a new multipart form based on string based fields.
func Fields(fields DataFields) p.Plugin {
return p.NewRequestPlugin(func(ctx *c.Context, h c.Handler) {
data := FormData{Data: fields}
handle(ctx, h, data)
})
}
// Data creates custom form based on the given form data
// who can have files and string based fields.
func Data(data FormData) p.Plugin {
return p.NewRequestPlugin(func(ctx *c.Context, h c.Handler) {
handle(ctx, h, data)
})
}
func handle(ctx *c.Context, h c.Handler, data FormData) {
if err := createForm(data, ctx); err != nil {
h.Error(ctx, err)
return
}
h.Next(ctx)
}
func createForm(data FormData, ctx *c.Context) error {
body := &bytes.Buffer{}
multipartWriter := multipart.NewWriter(body)
for index, file := range data.Files {
if err := writeFile(multipartWriter, data, file, index); err != nil {
return err
}
}
// Populate the other parts of the form (if there are any)
for key, values := range data.Data {
for _, value := range values {
multipartWriter.WriteField(key, value)
}
}
if err := multipartWriter.Close(); err != nil {
return err
}
ctx.Request.Method = setMethod(ctx)
ctx.Request.Body = ctx.WrapBody(ioutil.NopCloser(body))
ctx.Request.Header.Add("Content-Type", multipartWriter.FormDataContentType())
return nil
}
func writeFile(multipartWriter *multipart.Writer, data FormData, file FormFile, index int) error {
if file.Reader == nil {
return errors.New("gentleman: file reader cannot be nil")
}
rc, ok := file.Reader.(io.ReadCloser)
if !ok && file.Reader != nil {
rc = ioutil.NopCloser(file.Reader)
}
fileName := "file"
if len(data.Files) > 1 {
fileName = strings.Join([]string{fileName, strconv.Itoa(index + 1)}, "")
}
if file.Name != "" {
fileName = file.Name
}
writer, err := multipartWriter.CreateFormFile(fileName, file.Name)
if err != nil {
return err
}
if _, err = io.Copy(writer, rc); err != nil && err != io.EOF {
return err
}
rc.Close()
return nil
}
func setMethod(ctx *c.Context) string {
method := ctx.Request.Method
if method == "GET" || method == "" {
return "POST"
}
return method
}