Improve tls writer

transport/v2raywebsocket/client.go

@@ -74,7 +74,7 @@ func (c *Client) DialContext(ctx context.Context) (net.Conn, error) {
 	if c.maxEarlyData <= 0 {
 		conn, response, err := c.dialer.DialContext(ctx, c.uri, c.headers)
 		if err == nil {
-			return &WebsocketConn{Conn: conn, Writer: &Writer{conn, false}}, nil
+			return &WebsocketConn{Conn: conn, Writer: NewWriter(conn, false)}, nil
 		}
 		return nil, wrapDialError(response, err)
 	} else {

transport/v2raywebsocket/conn.go

@@ -26,7 +26,7 @@ func NewServerConn(wsConn *websocket.Conn, remoteAddr net.Addr) *WebsocketConn {
 	return &WebsocketConn{
 		Conn:       wsConn,
 		remoteAddr: remoteAddr,
-		Writer:     &Writer{wsConn, true},
+		Writer:     NewWriter(wsConn, true),
 	}
 }
 
@@ -117,7 +117,7 @@ func (c *EarlyWebsocketConn) Write(b []byte) (n int, err error) {
 	if err != nil {
 		return 0, wrapDialError(response, err)
 	}
-	c.conn = &WebsocketConn{Conn: conn, Writer: &Writer{conn, false}}
+	c.conn = &WebsocketConn{Conn: conn, Writer: NewWriter(conn, false)}
 	close(c.create)
 	if len(lateData) > 0 {
 		_, err = c.conn.Write(lateData)
@@ -160,7 +160,7 @@ func (c *EarlyWebsocketConn) WriteBuffer(buffer *buf.Buffer) error {
 	if err != nil {
 		return wrapDialError(response, err)
 	}
-	c.conn = &WebsocketConn{Conn: conn, Writer: &Writer{conn, false}}
+	c.conn = &WebsocketConn{Conn: conn, Writer: NewWriter(conn, false)}
 	close(c.create)
 	if len(lateData) > 0 {
 		_, err = c.conn.Write(lateData)

transport/v2raywebsocket/writer.go

@@ -4,8 +4,9 @@ import (
 	"encoding/binary"
 	"math/rand"
 
-	"github.com/sagernet/sing/common"
 	"github.com/sagernet/sing/common/buf"
+	"github.com/sagernet/sing/common/bufio"
+	N "github.com/sagernet/sing/common/network"
 	"github.com/sagernet/websocket"
 )
 
@@ -13,9 +14,18 @@ const frontHeadroom = 14
 
 type Writer struct {
 	*websocket.Conn
+	writer   N.ExtendedWriter
 	isServer bool
 }
 
+func NewWriter(conn *websocket.Conn, isServer bool) *Writer {
+	return &Writer{
+		conn,
+		bufio.NewExtendedWriter(conn.NetConn()),
+		isServer,
+	}
+}
+
 func (w *Writer) Write(p []byte) (n int, err error) {
 	err = w.Conn.WriteMessage(websocket.BinaryMessage, p)
 	if err != nil {
@@ -25,8 +35,6 @@ func (w *Writer) Write(p []byte) (n int, err error) {
 }
 
 func (w *Writer) WriteBuffer(buffer *buf.Buffer) error {
-	defer buffer.Release()
-
 	var payloadBitLength int
 	dataLen := buffer.Len()
 	data := buffer.Bytes()
@@ -69,5 +77,13 @@ func (w *Writer) WriteBuffer(buffer *buf.Buffer) error {
 		maskBytes(*(*[4]byte)(header[1+payloadBitLength:]), 0, data)
 	}
 
-	return common.Error(w.Conn.NetConn().Write(buffer.Bytes()))
+	return w.writer.WriteBuffer(buffer)
+}
+
+func (w *Writer) Upstream() any {
+	return w.Conn.NetConn()
+}
+
+func (w *Writer) FrontHeadroom() int {
+	return frontHeadroom
 }