github.com/inazumav/sing-box@v0.0.0-20230926072359-ab51429a14f1/transport/v2raywebsocket/writer.go (about)

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