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 func removeChunkExtension(p []byte) ([]byte, error) {
171 p, _, _ = bytes.Cut(p, semi)
172
173
174
175 return p, nil
176 }
177
178
179
180
181
182
183
184
185
186
187
188
189 func NewChunkedWriter(w io.Writer) io.WriteCloser {
190 return &chunkedWriter{w}
191 }
192
193
194
195 type chunkedWriter struct {
196 Wire io.Writer
197 }
198
199
200
201
202 func (cw *chunkedWriter) Write(data []byte) (n int, err error) {
203
204
205 if len(data) == 0 {
206 return 0, nil
207 }
208
209 if _, err = fmt.Fprintf(cw.Wire, "%x\r\n", len(data)); err != nil {
210 return 0, err
211 }
212 if n, err = cw.Wire.Write(data); err != nil {
213 return
214 }
215 if n != len(data) {
216 err = io.ErrShortWrite
217 return
218 }
219 if _, err = io.WriteString(cw.Wire, "\r\n"); err != nil {
220 return
221 }
222 if bw, ok := cw.Wire.(*FlushAfterChunkWriter); ok {
223 err = bw.Flush()
224 }
225 return
226 }
227
228 func (cw *chunkedWriter) Close() error {
229 _, err := io.WriteString(cw.Wire, "0\r\n")
230 return err
231 }
232
233
234
235
236
237
238 type FlushAfterChunkWriter struct {
239 *bufio.Writer
240 }
241
242 func parseHexUint(v []byte) (n uint64, err error) {
243 for i, b := range v {
244 switch {
245 case '0' <= b && b <= '9':
246 b = b - '0'
247 case 'a' <= b && b <= 'f':
248 b = b - 'a' + 10
249 case 'A' <= b && b <= 'F':
250 b = b - 'A' + 10
251 default:
252 return 0, errors.New("invalid byte in chunk length")
253 }
254 if i == 16 {
255 return 0, errors.New("http chunk length too large")
256 }
257 n <<= 4
258 n |= uint64(b)
259 }
260 return
261 }
262
View as plain text