-
Notifications
You must be signed in to change notification settings - Fork 4.6k
/
writer.go
126 lines (107 loc) · 2.64 KB
/
writer.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
package mux
import (
"v2ray.com/core/common"
"v2ray.com/core/common/buf"
"v2ray.com/core/common/net"
"v2ray.com/core/common/protocol"
"v2ray.com/core/common/serial"
)
type Writer struct {
dest net.Destination
writer buf.Writer
id uint16
followup bool
hasError bool
transferType protocol.TransferType
}
func NewWriter(id uint16, dest net.Destination, writer buf.Writer, transferType protocol.TransferType) *Writer {
return &Writer{
id: id,
dest: dest,
writer: writer,
followup: false,
transferType: transferType,
}
}
func NewResponseWriter(id uint16, writer buf.Writer, transferType protocol.TransferType) *Writer {
return &Writer{
id: id,
writer: writer,
followup: true,
transferType: transferType,
}
}
func (w *Writer) getNextFrameMeta() FrameMetadata {
meta := FrameMetadata{
SessionID: w.id,
Target: w.dest,
}
if w.followup {
meta.SessionStatus = SessionStatusKeep
} else {
w.followup = true
meta.SessionStatus = SessionStatusNew
}
return meta
}
func (w *Writer) writeMetaOnly() error {
meta := w.getNextFrameMeta()
b := buf.New()
if err := meta.WriteTo(b); err != nil {
return err
}
return w.writer.WriteMultiBuffer(buf.MultiBuffer{b})
}
func writeMetaWithFrame(writer buf.Writer, meta FrameMetadata, data buf.MultiBuffer) error {
frame := buf.New()
if err := meta.WriteTo(frame); err != nil {
return err
}
if _, err := serial.WriteUint16(frame, uint16(data.Len())); err != nil {
return err
}
mb2 := make(buf.MultiBuffer, 0, len(data)+1)
mb2 = append(mb2, frame)
mb2 = append(mb2, data...)
return writer.WriteMultiBuffer(mb2)
}
func (w *Writer) writeData(mb buf.MultiBuffer) error {
meta := w.getNextFrameMeta()
meta.Option.Set(OptionData)
return writeMetaWithFrame(w.writer, meta, mb)
}
// WriteMultiBuffer implements buf.Writer.
func (w *Writer) WriteMultiBuffer(mb buf.MultiBuffer) error {
defer buf.ReleaseMulti(mb)
if mb.IsEmpty() {
return w.writeMetaOnly()
}
for !mb.IsEmpty() {
var chunk buf.MultiBuffer
if w.transferType == protocol.TransferTypeStream {
mb, chunk = buf.SplitSize(mb, 8*1024)
} else {
mb2, b := buf.SplitFirst(mb)
mb = mb2
chunk = buf.MultiBuffer{b}
}
if err := w.writeData(chunk); err != nil {
return err
}
}
return nil
}
// Close implements common.Closable.
func (w *Writer) Close() error {
meta := FrameMetadata{
SessionID: w.id,
SessionStatus: SessionStatusEnd,
}
if w.hasError {
meta.Option.Set(OptionError)
}
frame := buf.New()
common.Must(meta.WriteTo(frame))
w.writer.WriteMultiBuffer(buf.MultiBuffer{frame}) // nolint: errcheck
return nil
}