-
Notifications
You must be signed in to change notification settings - Fork 0
/
request_reader.go
200 lines (174 loc) · 4.19 KB
/
request_reader.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
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
package go_readers
import (
"bytes"
"crypto/rand"
"fmt"
"io"
"net/textproto"
"sort"
"strings"
)
const (
partitionStateNew = iota
partitionStateHeaderSent
partitionStateSent
)
var (
ErrPartitionReaderIsNil = fmt.Errorf("partition header or body is nil")
ErrPartitionInvalidState = fmt.Errorf("invalid partition state")
ErrInvalidReadSize = fmt.Errorf("invalid read size")
)
type MultipartRequestBody interface {
io.ReadCloser
Boundary() string
NewPart(header textproto.MIMEHeader, sz int64, reader io.ReadCloser)
Size() int64
}
func NewMultipartRequestBody() MultipartRequestBody {
var buf [30]byte
_, _ = io.ReadFull(rand.Reader, buf[:])
boundary := fmt.Sprintf("%x", buf[:])
return &multipartRequestBody{
boundary: boundary,
parts: readInfo{},
partList: make([]*partition, 0),
closeReader: strings.NewReader(fmt.Sprintf("\r\n--%s--\r\n", boundary)),
}
}
type multipartRequestBody struct {
boundary string
parts readInfo
partList []*partition
closeReader io.Reader
cntRead int
}
func (inst *multipartRequestBody) NewPart(header textproto.MIMEHeader, sz int64, reader io.ReadCloser) {
var b bytes.Buffer
keys := make([]string, 0, len(header))
for k := range header {
keys = append(keys, k)
}
sort.Strings(keys)
for _, k := range keys {
for _, v := range header[k] {
_, _ = fmt.Fprintf(&b, "%s: %s\r\n", k, v)
}
}
_, _ = fmt.Fprintf(&b, "\r\n")
inst.partList = append(inst.partList, &partition{
&b,
reader,
int64(b.Len()),
sz,
0,
partitionStateNew,
})
}
func (inst *multipartRequestBody) Read(p []byte) (n int, err error) {
return inst.read(p)
}
func (inst *multipartRequestBody) read(p []byte) (n int, err error) {
if inst.parts.partIdx < len(inst.partList) { // has unread partitions
if inst.parts.boundaryIdx == inst.parts.partIdx { // boundary is not read
if inst.parts.boundary == nil {
inst.parts.boundary = strings.NewReader(fmt.Sprintf("\r\n--%s\r\n", inst.boundary))
}
n, err = inst.parts.boundary.Read(p)
if err == io.EOF { // boundary is read
inst.parts.boundaryIdx++
inst.parts.boundary = nil
return inst.read(p)
}
} else { // partition is not read
n, err = inst.partList[inst.parts.partIdx].Read(p)
if err == io.EOF { // partition is read
inst.parts.partIdx++
return inst.read(p)
}
}
inst.cntRead += n
return
}
n, err = inst.closeReader.Read(p)
if err == nil && int64(n+inst.cntRead) != inst.Size() {
return 0, fmt.Errorf("read %d != size %d: %w", n+inst.cntRead, inst.Size(), ErrInvalidReadSize)
}
return
}
func (inst *multipartRequestBody) Close() (err error) {
for _, r := range inst.partList {
if lErr := r.Close(); lErr != nil {
err = lErr
}
}
return
}
func (inst *multipartRequestBody) Size() (sz int64) {
// boundary size
szBoundary := int64(len(inst.boundary) + 6)
for _, r := range inst.partList {
sz += r.Size() + szBoundary
}
// last boundary size
szBoundary += 2
return sz + szBoundary
}
func (inst *multipartRequestBody) Boundary() string {
return inst.boundary
}
type partition struct {
header io.Reader
reader io.ReadCloser
szHeader int64
szReader int64
cntRead int
state int
}
func (inst *partition) Read(p []byte) (n int, err error) {
if inst.header == nil || inst.reader == nil {
return 0, ErrPartitionReaderIsNil
}
switch inst.state {
case partitionStateNew:
n, err = inst.header.Read(p)
if err == io.EOF {
if n > 0 {
err = nil
} else {
inst.state = partitionStateHeaderSent
return inst.Read(p)
}
}
inst.cntRead += n
return
case partitionStateHeaderSent:
n, err = inst.reader.Read(p)
if err == io.EOF {
if n > 0 {
err = nil
} else {
inst.state = partitionStateSent
return inst.Read(p)
}
}
inst.cntRead += n
return
case partitionStateSent:
if int64(inst.cntRead) != inst.Size() {
return 0, fmt.Errorf("read %d != size %d: %w", inst.cntRead, inst.Size(), ErrInvalidReadSize)
}
return 0, io.EOF
}
return 0, ErrPartitionInvalidState
}
func (inst *partition) Close() error {
return inst.reader.Close()
}
func (inst *partition) Size() int64 {
return inst.szHeader + inst.szReader
}
type readInfo struct {
partIdx int
boundaryIdx int
boundary io.Reader
}