cleanup

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))
 }