diff --git a/app/proxyman/mux/writer.go b/app/proxyman/mux/writer.go index ca3cea654..800400301 100644 --- a/app/proxyman/mux/writer.go +++ b/app/proxyman/mux/writer.go @@ -1,6 +1,8 @@ package mux import ( + "runtime" + "v2ray.com/core/common/buf" "v2ray.com/core/common/net" "v2ray.com/core/common/serial" @@ -30,7 +32,7 @@ func NewResponseWriter(id uint16, writer buf.Writer) *Writer { } } -func (w *Writer) writeInternal(mb buf.MultiBuffer) error { +func (w *Writer) getNextFrameMeta() FrameMetadata { meta := FrameMetadata{ SessionID: w.id, Target: w.dest, @@ -42,35 +44,51 @@ func (w *Writer) writeInternal(mb buf.MultiBuffer) error { meta.SessionStatus = SessionStatusNew } - hasData := !mb.IsEmpty() + return meta +} - if hasData { - meta.Option.Add(OptionData) +func (w *Writer) writeMetaOnly() error { + meta := w.getNextFrameMeta() + b := buf.New() + if err := b.AppendSupplier(meta.AsSupplier()); err != nil { + return err } + runtime.KeepAlive(meta) + return w.writer.Write(buf.NewMultiBufferValue(b)) +} + +func (w *Writer) writeData(mb buf.MultiBuffer) error { + meta := w.getNextFrameMeta() + meta.Option.Add(OptionData) frame := buf.New() - frame.AppendSupplier(meta.AsSupplier()) + if err := frame.AppendSupplier(meta.AsSupplier()); err != nil { + return err + } + runtime.KeepAlive(meta) mb2 := buf.NewMultiBuffer() mb2.Append(frame) - if hasData { - frame.AppendSupplier(serial.WriteUint16(uint16(mb.Len()))) - mb2.AppendMulti(mb) + if err := frame.AppendSupplier(serial.WriteUint16(uint16(mb.Len()))); err != nil { + return err } + + mb2.AppendMulti(mb) return w.writer.Write(mb2) } func (w *Writer) Write(mb buf.MultiBuffer) error { + if mb.IsEmpty() { + return w.writeMetaOnly() + } + const chunkSize = 8 * 1024 - for { + for !mb.IsEmpty() { slice := mb.SliceBySize(chunkSize) - if err := w.writeInternal(slice); err != nil { + if err := w.writeData(slice); err != nil { return err } - if mb.IsEmpty() { - break - } } return nil } @@ -83,6 +101,7 @@ func (w *Writer) Close() { frame := buf.New() frame.AppendSupplier(meta.AsSupplier()) + runtime.KeepAlive(meta) w.writer.Write(buf.NewMultiBufferValue(frame)) }