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
|
package mux
import (
"v2ray.com/core/common"
"v2ray.com/core/common/buf"
"v2ray.com/core/common/errors"
"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
}
if len(data)+1 > 64*1024*1024 {
return errors.New("value too large")
}
sliceSize := len(data) + 1
mb2 := make(buf.MultiBuffer, 0, sliceSize)
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})
return nil
}
|