github.com/sagernet/sing-box@v1.9.0-rc.20/transport/v2raywebsocket/writer.go (about)

     1  package v2raywebsocket
     2  
     3  import (
     4  	"encoding/binary"
     5  	"io"
     6  	"math/rand"
     7  
     8  	"github.com/sagernet/sing/common/buf"
     9  	"github.com/sagernet/sing/common/bufio"
    10  	N "github.com/sagernet/sing/common/network"
    11  	"github.com/sagernet/ws"
    12  )
    13  
    14  type Writer struct {
    15  	writer   N.ExtendedWriter
    16  	isServer bool
    17  }
    18  
    19  func NewWriter(writer io.Writer, state ws.State) *Writer {
    20  	return &Writer{
    21  		bufio.NewExtendedWriter(writer),
    22  		state == ws.StateServerSide,
    23  	}
    24  }
    25  
    26  func (w *Writer) WriteBuffer(buffer *buf.Buffer) error {
    27  	var payloadBitLength int
    28  	dataLen := buffer.Len()
    29  	data := buffer.Bytes()
    30  	if dataLen < 126 {
    31  		payloadBitLength = 1
    32  	} else if dataLen < 65536 {
    33  		payloadBitLength = 3
    34  	} else {
    35  		payloadBitLength = 9
    36  	}
    37  
    38  	var headerLen int
    39  	headerLen += 1 // FIN / RSV / OPCODE
    40  	headerLen += payloadBitLength
    41  	if !w.isServer {
    42  		headerLen += 4 // MASK KEY
    43  	}
    44  
    45  	header := buffer.ExtendHeader(headerLen)
    46  	header[0] = byte(ws.OpBinary) | 0x80
    47  	if w.isServer {
    48  		header[1] = 0
    49  	} else {
    50  		header[1] = 1 << 7
    51  	}
    52  
    53  	if dataLen < 126 {
    54  		header[1] |= byte(dataLen)
    55  	} else if dataLen < 65536 {
    56  		header[1] |= 126
    57  		binary.BigEndian.PutUint16(header[2:], uint16(dataLen))
    58  	} else {
    59  		header[1] |= 127
    60  		binary.BigEndian.PutUint64(header[2:], uint64(dataLen))
    61  	}
    62  
    63  	if !w.isServer {
    64  		maskKey := rand.Uint32()
    65  		binary.BigEndian.PutUint32(header[1+payloadBitLength:], maskKey)
    66  		ws.Cipher(data, *(*[4]byte)(header[1+payloadBitLength:]), 0)
    67  	}
    68  
    69  	return w.writer.WriteBuffer(buffer)
    70  }
    71  
    72  func (w *Writer) FrontHeadroom() int {
    73  	return 14
    74  }