v2ray-core/app/proxyman/mux/writer.go

package mux

import (
	"runtime"

	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/serial"
)

type Writer struct {
	id       uint16
	dest     net.Destination
	writer   buf.Writer
	followup bool
}

func NewWriter(id uint16, dest net.Destination, writer buf.Writer) *Writer {
	return &Writer{
		id:       id,
		dest:     dest,
		writer:   writer,
		followup: false,
	}
}

func NewResponseWriter(id uint16, writer buf.Writer) *Writer {
	return &Writer{
		id:       id,
		writer:   writer,
		followup: true,
	}
}

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 := 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()
	if err := frame.AppendSupplier(meta.AsSupplier()); err != nil {
		return err
	}
	runtime.KeepAlive(meta)
	if err := frame.AppendSupplier(serial.WriteUint16(uint16(mb.Len()))); err != nil {
		return err
	}

	mb2 := buf.NewMultiBuffer()
	mb2.Append(frame)
	mb2.AppendMulti(mb)
	return w.writer.Write(mb2)
}

// Write implements buf.MultiBufferWriter.
func (w *Writer) Write(mb buf.MultiBuffer) error {
	if mb.IsEmpty() {
		return w.writeMetaOnly()
	}

	const chunkSize = 8 * 1024
	for !mb.IsEmpty() {
		slice := mb.SliceBySize(chunkSize)
		if err := w.writeData(slice); err != nil {
			return err
		}
	}
	return nil
}

func (w *Writer) Close() {
	meta := FrameMetadata{
		SessionID:     w.id,
		SessionStatus: SessionStatusEnd,
	}

	frame := buf.New()
	frame.AppendSupplier(meta.AsSupplier())
	runtime.KeepAlive(meta)

	w.writer.Write(buf.NewMultiBufferValue(frame))
}