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

     1  package mux
     2  
     3  import (
     4  	"github.com/xraypb/xray-core/common"
     5  	"github.com/xraypb/xray-core/common/buf"
     6  	"github.com/xraypb/xray-core/common/net"
     7  	"github.com/xraypb/xray-core/common/protocol"
     8  	"github.com/xraypb/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  }
    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 len(data) == 1 {
    67  		frame.UDP = data[0].UDP
    68  	}
    69  	if err := meta.WriteTo(frame); err != nil {
    70  		return err
    71  	}
    72  	if _, err := serial.WriteUint16(frame, uint16(data.Len())); err != nil {
    73  		return err
    74  	}
    75  
    76  	mb2 := make(buf.MultiBuffer, 0, len(data)+1)
    77  	mb2 = append(mb2, frame)
    78  	mb2 = append(mb2, data...)
    79  	return writer.WriteMultiBuffer(mb2)
    80  }
    81  
    82  func (w *Writer) writeData(mb buf.MultiBuffer) error {
    83  	meta := w.getNextFrameMeta()
    84  	meta.Option.Set(OptionData)
    85  
    86  	return writeMetaWithFrame(w.writer, meta, mb)
    87  }
    88  
    89  // WriteMultiBuffer implements buf.Writer.
    90  func (w *Writer) WriteMultiBuffer(mb buf.MultiBuffer) error {
    91  	defer buf.ReleaseMulti(mb)
    92  
    93  	if mb.IsEmpty() {
    94  		return w.writeMetaOnly()
    95  	}
    96  
    97  	for !mb.IsEmpty() {
    98  		var chunk buf.MultiBuffer
    99  		if w.transferType == protocol.TransferTypeStream {
   100  			mb, chunk = buf.SplitSize(mb, 8*1024)
   101  		} else {
   102  			mb2, b := buf.SplitFirst(mb)
   103  			mb = mb2
   104  			chunk = buf.MultiBuffer{b}
   105  		}
   106  		if err := w.writeData(chunk); err != nil {
   107  			return err
   108  		}
   109  	}
   110  
   111  	return nil
   112  }
   113  
   114  // Close implements common.Closable.
   115  func (w *Writer) Close() error {
   116  	meta := FrameMetadata{
   117  		SessionID:     w.id,
   118  		SessionStatus: SessionStatusEnd,
   119  	}
   120  	if w.hasError {
   121  		meta.Option.Set(OptionError)
   122  	}
   123  
   124  	frame := buf.New()
   125  	common.Must(meta.WriteTo(frame))
   126  
   127  	w.writer.WriteMultiBuffer(buf.MultiBuffer{frame})
   128  	return nil
   129  }