1
2
3
4
5
6
7
8
9 package internal
10
11 import (
12 "bufio"
13 "bytes"
14 "errors"
15 "fmt"
16 "io"
17 )
18
19 const maxLineLength = 4096
20
21 var ErrLineTooLong = errors.New("header line too long")
22
23
24
25
26
27
28
29 func NewChunkedReader(r io.Reader) io.Reader {
30 br, ok := r.(*bufio.Reader)
31 if !ok {
32 br = bufio.NewReader(r)
33 }
34 return &chunkedReader{r: br}
35 }
36
37 type chunkedReader struct {
38 r *bufio.Reader
39 n uint64
40 err error
41 buf [2]byte
42 checkEnd bool
43 }
44
45 func (cr *chunkedReader) beginChunk() {
46
47 var line []byte
48 line, cr.err = readChunkLine(cr.r)
49 if cr.err != nil {
50 return
51 }
52 cr.n, cr.err = parseHexUint(line)
53 if cr.err != nil {
54 return
55 }
56 if cr.n == 0 {
57 cr.err = io.EOF
58 }
59 }
60
61 func (cr *chunkedReader) chunkHeaderAvailable() bool {
62 n := cr.r.Buffered()
63 if n > 0 {
64 peek, _ := cr.r.Peek(n)
65 return bytes.IndexByte(peek, '\n') >= 0
66 }
67 return false
68 }
69
70 func (cr *chunkedReader) Read(b []uint8) (n int, err error) {
71 for cr.err == nil {
72 if cr.checkEnd {
73 if n > 0 && cr.r.Buffered() < 2 {
74
75
76
77 break
78 }
79 if _, cr.err = io.ReadFull(cr.r, cr.buf[:2]); cr.err == nil {
80 if string(cr.buf[:]) != "\r\n" {
81 cr.err = errors.New("malformed chunked encoding")
82 break
83 }
84 } else {
85 if cr.err == io.EOF {
86 cr.err = io.ErrUnexpectedEOF
87 }
88 break
89 }
90 cr.checkEnd = false
91 }
92 if cr.n == 0 {
93 if n > 0 && !cr.chunkHeaderAvailable() {
94
95
96 break
97 }
98 cr.beginChunk()
99 continue
100 }
101 if len(b) == 0 {
102 break
103 }
104 rbuf := b
105 if uint64(len(rbuf)) > cr.n {
106 rbuf = rbuf[:cr.n]
107 }
108 var n0 int
109 n0, cr.err = cr.r.Read(rbuf)
110 n += n0
111 b = b[n0:]
112 cr.n -= uint64(n0)
113
114
115 if cr.n == 0 && cr.err == nil {
116 cr.checkEnd = true
117 } else if cr.err == io.EOF {
118 cr.err = io.ErrUnexpectedEOF
119 }
120 }
121 return n, cr.err
122 }
123
124
125
126
127
128 func readChunkLine(b *bufio.Reader) ([]byte, error) {
129 p, err := b.ReadSlice('\n')
130 if err != nil {
131
132
133 if err == io.EOF {
134 err = io.ErrUnexpectedEOF
135 } else if err == bufio.ErrBufferFull {
136 err = ErrLineTooLong
137 }
138 return nil, err
139 }
140 if len(p) >= maxLineLength {
141 return nil, ErrLineTooLong
142 }
143 p = trimTrailingWhitespace(p)
144 p, err = removeChunkExtension(p)
145 if err != nil {
146 return nil, err
147 }
148 return p, nil
149 }
150
151 func trimTrailingWhitespace(b []byte) []byte {
152 for len(b) > 0 && isASCIISpace(b[len(b)-1]) {
153 b = b[:len(b)-1]
154 }
155 return b
156 }
157
158 func isASCIISpace(b byte) bool {
159 return b == ' ' || b == '\t' || b == '\n' || b == '\r'
160 }
161
162 var semi = []byte(";")
163
164
165
166
167
168
169
170
171 func removeChunkExtension(p []byte) ([]byte, error) {
172 p, _, _ = bytes.Cut(p, semi)
173
174
175
176 return p, nil
177 }
178
179
180
181
182
183
184
185
186
187
188
189
190 func NewChunkedWriter(w io.Writer) io.WriteCloser {
191 return &chunkedWriter{w}
192 }
193
194
195
196 type chunkedWriter struct {
197 Wire io.Writer
198 }
199
200
201
202
203 func (cw *chunkedWriter) Write(data []byte) (n int, err error) {
204
205
206 if len(data) == 0 {
207 return 0, nil
208 }
209
210 if _, err = fmt.Fprintf(cw.Wire, "%x\r\n", len(data)); err != nil {
211 return 0, err
212 }
213 if n, err = cw.Wire.Write(data); err != nil {
214 return
215 }
216 if n != len(data) {
217 err = io.ErrShortWrite
218 return
219 }
220 if _, err = io.WriteString(cw.Wire, "\r\n"); err != nil {
221 return
222 }
223 if bw, ok := cw.Wire.(*FlushAfterChunkWriter); ok {
224 err = bw.Flush()
225 }
226 return
227 }
228
229 func (cw *chunkedWriter) Close() error {
230 _, err := io.WriteString(cw.Wire, "0\r\n")
231 return err
232 }
233
234
235
236
237
238
239 type FlushAfterChunkWriter struct {
240 *bufio.Writer
241 }
242
243 func parseHexUint(v []byte) (n uint64, err error) {
244 for i, b := range v {
245 switch {
246 case '0' <= b && b <= '9':
247 b = b - '0'
248 case 'a' <= b && b <= 'f':
249 b = b - 'a' + 10
250 case 'A' <= b && b <= 'F':
251 b = b - 'A' + 10
252 default:
253 return 0, errors.New("invalid byte in chunk length")
254 }
255 if i == 16 {
256 return 0, errors.New("http chunk length too large")
257 }
258 n <<= 4
259 n |= uint64(b)
260 }
261 return
262 }
263
View as plain text