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
|
package fasthttp
import (
"bufio"
"bytes"
"io"
"sync"
"github.com/valyala/bytebufferpool"
)
type headerInterface interface {
ContentLength() int
ReadTrailer(r *bufio.Reader) error
}
type requestStream struct {
header headerInterface
prefetchedBytes *bytes.Reader
reader *bufio.Reader
totalBytesRead int
chunkLeft int
}
func (rs *requestStream) Read(p []byte) (int, error) {
var (
n int
err error
)
if rs.header.ContentLength() == -1 {
if rs.chunkLeft == 0 {
chunkSize, err := parseChunkSize(rs.reader)
if err != nil {
return 0, err
}
if chunkSize == 0 {
err = rs.header.ReadTrailer(rs.reader)
if err != nil && err != io.EOF {
return 0, err
}
return 0, io.EOF
}
rs.chunkLeft = chunkSize
}
bytesToRead := len(p)
if rs.chunkLeft < len(p) {
bytesToRead = rs.chunkLeft
}
n, err = rs.reader.Read(p[:bytesToRead])
rs.totalBytesRead += n
rs.chunkLeft -= n
if err == io.EOF {
err = io.ErrUnexpectedEOF
}
if err == nil && rs.chunkLeft == 0 {
err = readCrLf(rs.reader)
}
return n, err
}
if rs.totalBytesRead == rs.header.ContentLength() {
return 0, io.EOF
}
prefetchedSize := int(rs.prefetchedBytes.Size())
if prefetchedSize > rs.totalBytesRead {
left := prefetchedSize - rs.totalBytesRead
if len(p) > left {
p = p[:left]
}
n, err := rs.prefetchedBytes.Read(p)
rs.totalBytesRead += n
if n == rs.header.ContentLength() {
return n, io.EOF
}
return n, err
}
left := rs.header.ContentLength() - rs.totalBytesRead
if len(p) > left {
p = p[:left]
}
n, err = rs.reader.Read(p)
rs.totalBytesRead += n
if err != nil {
return n, err
}
if rs.totalBytesRead == rs.header.ContentLength() {
err = io.EOF
}
return n, err
}
func acquireRequestStream(b *bytebufferpool.ByteBuffer, r *bufio.Reader, h headerInterface) *requestStream {
rs := requestStreamPool.Get().(*requestStream)
rs.prefetchedBytes = bytes.NewReader(b.B)
rs.reader = r
rs.header = h
return rs
}
func releaseRequestStream(rs *requestStream) {
rs.prefetchedBytes = nil
rs.totalBytesRead = 0
rs.chunkLeft = 0
rs.reader = nil
rs.header = nil
requestStreamPool.Put(rs)
}
var requestStreamPool = sync.Pool{
New: func() any {
return &requestStream{}
},
}
|
