github.com/xmplusdev/xray-core@v1.8.10/common/mux/writer.go (about)

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