github.com/imannamdari/v2ray-core/v5@v5.0.5/common/mux/writer.go (about)

     1  package mux
     2  
     3  import (
     4  	"github.com/imannamdari/v2ray-core/v5/common"
     5  	"github.com/imannamdari/v2ray-core/v5/common/buf"
     6  	"github.com/imannamdari/v2ray-core/v5/common/errors"
     7  	"github.com/imannamdari/v2ray-core/v5/common/net"
     8  	"github.com/imannamdari/v2ray-core/v5/common/protocol"
     9  	"github.com/imannamdari/v2ray-core/v5/common/serial"
    10  )
    11  
    12  type Writer struct {
    13  	dest         net.Destination
    14  	writer       buf.Writer
    15  	id           uint16
    16  	followup     bool
    17  	hasError     bool
    18  	transferType protocol.TransferType
    19  }
    20  
    21  func NewWriter(id uint16, dest net.Destination, writer buf.Writer, transferType protocol.TransferType) *Writer {
    22  	return &Writer{
    23  		id:           id,
    24  		dest:         dest,
    25  		writer:       writer,
    26  		followup:     false,
    27  		transferType: transferType,
    28  	}
    29  }
    30  
    31  func NewResponseWriter(id uint16, writer buf.Writer, transferType protocol.TransferType) *Writer {
    32  	return &Writer{
    33  		id:           id,
    34  		writer:       writer,
    35  		followup:     true,
    36  		transferType: transferType,
    37  	}
    38  }
    39  
    40  func (w *Writer) getNextFrameMeta() FrameMetadata {
    41  	meta := FrameMetadata{
    42  		SessionID: w.id,
    43  		Target:    w.dest,
    44  	}
    45  
    46  	if w.followup {
    47  		meta.SessionStatus = SessionStatusKeep
    48  	} else {
    49  		w.followup = true
    50  		meta.SessionStatus = SessionStatusNew
    51  	}
    52  
    53  	return meta
    54  }
    55  
    56  func (w *Writer) writeMetaOnly() error {
    57  	meta := w.getNextFrameMeta()
    58  	b := buf.New()
    59  	if err := meta.WriteTo(b); err != nil {
    60  		return err
    61  	}
    62  	return w.writer.WriteMultiBuffer(buf.MultiBuffer{b})
    63  }
    64  
    65  func writeMetaWithFrame(writer buf.Writer, meta FrameMetadata, data buf.MultiBuffer) error {
    66  	frame := buf.New()
    67  	if err := meta.WriteTo(frame); err != nil {
    68  		return err
    69  	}
    70  	if _, err := serial.WriteUint16(frame, uint16(data.Len())); err != nil {
    71  		return err
    72  	}
    73  
    74  	if len(data)+1 > 64*1024*1024 {
    75  		return errors.New("value too large")
    76  	}
    77  	sliceSize := len(data) + 1
    78  	mb2 := make(buf.MultiBuffer, 0, sliceSize)
    79  	mb2 = append(mb2, frame)
    80  	mb2 = append(mb2, data...)
    81  	return writer.WriteMultiBuffer(mb2)
    82  }
    83  
    84  func (w *Writer) writeData(mb buf.MultiBuffer) error {
    85  	meta := w.getNextFrameMeta()
    86  	meta.Option.Set(OptionData)
    87  
    88  	return writeMetaWithFrame(w.writer, meta, mb)
    89  }
    90  
    91  // WriteMultiBuffer implements buf.Writer.
    92  func (w *Writer) WriteMultiBuffer(mb buf.MultiBuffer) error {
    93  	defer buf.ReleaseMulti(mb)
    94  
    95  	if mb.IsEmpty() {
    96  		return w.writeMetaOnly()
    97  	}
    98  
    99  	for !mb.IsEmpty() {
   100  		var chunk buf.MultiBuffer
   101  		if w.transferType == protocol.TransferTypeStream {
   102  			mb, chunk = buf.SplitSize(mb, 8*1024)
   103  		} else {
   104  			mb2, b := buf.SplitFirst(mb)
   105  			mb = mb2
   106  			chunk = buf.MultiBuffer{b}
   107  		}
   108  		if err := w.writeData(chunk); err != nil {
   109  			return err
   110  		}
   111  	}
   112  
   113  	return nil
   114  }
   115  
   116  // Close implements common.Closable.
   117  func (w *Writer) Close() error {
   118  	meta := FrameMetadata{
   119  		SessionID:     w.id,
   120  		SessionStatus: SessionStatusEnd,
   121  	}
   122  	if w.hasError {
   123  		meta.Option.Set(OptionError)
   124  	}
   125  
   126  	frame := buf.New()
   127  	common.Must(meta.WriteTo(frame))
   128  
   129  	w.writer.WriteMultiBuffer(buf.MultiBuffer{frame})
   130  	return nil
   131  }