forked from v2ray/v2ray-core
-
Notifications
You must be signed in to change notification settings - Fork 0
/
chunk.go
147 lines (121 loc) · 2.89 KB
/
chunk.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
package crypto
import (
"io"
"v2ray.com/core/common/buf"
"v2ray.com/core/common/serial"
)
type ChunkSizeDecoder interface {
SizeBytes() int
Decode([]byte) (uint16, error)
}
type ChunkSizeEncoder interface {
SizeBytes() int
Encode(uint16, []byte) []byte
}
type PlainChunkSizeParser struct{}
func (PlainChunkSizeParser) SizeBytes() int {
return 2
}
func (PlainChunkSizeParser) Encode(size uint16, b []byte) []byte {
return serial.Uint16ToBytes(size, b)
}
func (PlainChunkSizeParser) Decode(b []byte) (uint16, error) {
return serial.BytesToUint16(b), nil
}
type ChunkStreamReader struct {
sizeDecoder ChunkSizeDecoder
reader buf.Reader
buffer []byte
leftOver buf.MultiBuffer
leftOverSize int
}
func NewChunkStreamReader(sizeDecoder ChunkSizeDecoder, reader io.Reader) *ChunkStreamReader {
return &ChunkStreamReader{
sizeDecoder: sizeDecoder,
reader: buf.NewReader(reader),
buffer: make([]byte, sizeDecoder.SizeBytes()),
}
}
func (r *ChunkStreamReader) readAtLeast(size int) error {
mb := r.leftOver
r.leftOver = nil
for mb.Len() < size {
extra, err := r.reader.Read()
if err != nil {
mb.Release()
return err
}
mb.AppendMulti(extra)
}
r.leftOver = mb
return nil
}
func (r *ChunkStreamReader) readSize() (uint16, error) {
if r.sizeDecoder.SizeBytes() > r.leftOver.Len() {
if err := r.readAtLeast(r.sizeDecoder.SizeBytes() - r.leftOver.Len()); err != nil {
return 0, err
}
}
r.leftOver.Read(r.buffer)
return r.sizeDecoder.Decode(r.buffer)
}
func (r *ChunkStreamReader) Read() (buf.MultiBuffer, error) {
size := r.leftOverSize
if size == 0 {
nextSize, err := r.readSize()
if err != nil {
return nil, err
}
if nextSize == 0 {
return nil, io.EOF
}
size = int(nextSize)
}
if r.leftOver.IsEmpty() {
if err := r.readAtLeast(1); err != nil {
return nil, err
}
}
if size >= r.leftOver.Len() {
mb := r.leftOver
r.leftOverSize = size - r.leftOver.Len()
r.leftOver = nil
return mb, nil
}
mb := r.leftOver.SliceBySize(size)
if mb.Len() != size {
b := buf.New()
b.AppendSupplier(buf.ReadFullFrom(&r.leftOver, size-mb.Len()))
mb.Append(b)
}
r.leftOverSize = 0
return mb, nil
}
type ChunkStreamWriter struct {
sizeEncoder ChunkSizeEncoder
writer buf.Writer
}
func NewChunkStreamWriter(sizeEncoder ChunkSizeEncoder, writer io.Writer) *ChunkStreamWriter {
return &ChunkStreamWriter{
sizeEncoder: sizeEncoder,
writer: buf.NewWriter(writer),
}
}
func (w *ChunkStreamWriter) Write(mb buf.MultiBuffer) error {
mb2Write := buf.NewMultiBuffer()
const sliceSize = 8192
for {
slice := mb.SliceBySize(sliceSize)
b := buf.New()
b.AppendSupplier(func(buffer []byte) (int, error) {
w.sizeEncoder.Encode(uint16(slice.Len()), buffer[:0])
return w.sizeEncoder.SizeBytes(), nil
})
mb2Write.Append(b)
mb2Write.AppendMulti(slice)
if mb.IsEmpty() {
break
}
}
return w.writer.Write(mb2Write)
}