github.com/Uhtred009/v2ray-core-1@v4.31.2+incompatible/common/mux/writer.go (about)

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