Fix deadline usage

transport/v2raywebsocket/client.go

@@ -12,6 +12,9 @@ import (
 	"github.com/sagernet/sing-box/common/tls"
 	C "github.com/sagernet/sing-box/constant"
 	"github.com/sagernet/sing-box/option"
+	"github.com/sagernet/sing/common/buf"
+	"github.com/sagernet/sing/common/bufio"
+	"github.com/sagernet/sing/common/bufio/deadline"
 	E "github.com/sagernet/sing/common/exceptions"
 	M "github.com/sagernet/sing/common/metadata"
 	N "github.com/sagernet/sing/common/network"
@@ -87,18 +90,35 @@ func (c *Client) dialContext(ctx context.Context, requestURL *url.URL, headers h
 			return nil, err
 		}
 	}
-	conn.SetDeadline(time.Now().Add(C.TCPTimeout))
+	var deadlineConn net.Conn
+	if deadline.NeedAdditionalReadDeadline(conn) {
+		deadlineConn = deadline.NewConn(conn)
+	} else {
+		deadlineConn = conn
+	}
+	err = deadlineConn.SetDeadline(time.Now().Add(C.TCPTimeout))
+	if err != nil {
+		return nil, E.Cause(err, "set read deadline")
+	}
 	var protocols []string
 	if protocolHeader := headers.Get("Sec-WebSocket-Protocol"); protocolHeader != "" {
 		protocols = []string{protocolHeader}
 		headers.Del("Sec-WebSocket-Protocol")
 	}
-	reader, _, err := ws.Dialer{Header: ws.HandshakeHeaderHTTP(headers), Protocols: protocols}.Upgrade(conn, requestURL)
-	conn.SetDeadline(time.Time{})
+	reader, _, err := ws.Dialer{Header: ws.HandshakeHeaderHTTP(headers), Protocols: protocols}.Upgrade(deadlineConn, requestURL)
+	deadlineConn.SetDeadline(time.Time{})
 	if err != nil {
 		return nil, err
 	}
-	return NewConn(conn, reader, nil, ws.StateClientSide), nil
+	if reader != nil {
+		buffer := buf.NewSize(reader.Buffered())
+		_, err = buffer.ReadFullFrom(reader, buffer.Len())
+		if err != nil {
+			return nil, err
+		}
+		conn = bufio.NewCachedConn(conn, buffer)
+	}
+	return NewConn(conn, nil, ws.StateClientSide), nil
 }
 
 func (c *Client) DialContext(ctx context.Context) (net.Conn, error) {

transport/v2raywebsocket/conn.go

@@ -1,7 +1,6 @@
 package v2raywebsocket
 
 import (
-	"bufio"
 	"context"
 	"encoding/base64"
 	"io"
@@ -28,19 +27,13 @@ type WebsocketConn struct {
 	remoteAddr     net.Addr
 }
 
-func NewConn(conn net.Conn, br *bufio.Reader, remoteAddr net.Addr, state ws.State) *WebsocketConn {
+func NewConn(conn net.Conn, remoteAddr net.Addr, state ws.State) *WebsocketConn {
 	controlHandler := wsutil.ControlFrameHandler(conn, state)
-	var reader io.Reader
-	if br != nil && br.Buffered() > 0 {
-		reader = br
-	} else {
-		reader = conn
-	}
 	return &WebsocketConn{
 		Conn:  conn,
 		state: state,
 		reader: &wsutil.Reader{
-			Source:          reader,
+			Source:          conn,
 			State:           state,
 			SkipHeaderCheck: !debug.Enabled,
 			OnIntermediate:  controlHandler,

transport/v2raywebsocket/server.go

@@ -88,14 +88,14 @@ func (s *Server) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
 		s.invalidRequest(writer, request, http.StatusBadRequest, E.Cause(err, "decode early data"))
 		return
 	}
-	wsConn, reader, _, err := ws.UpgradeHTTP(request, writer)
+	wsConn, _, _, err := ws.UpgradeHTTP(request, writer)
 	if err != nil {
 		s.invalidRequest(writer, request, 0, E.Cause(err, "upgrade websocket connection"))
 		return
 	}
 	var metadata M.Metadata
 	metadata.Source = sHttp.SourceAddress(request)
-	conn = NewConn(wsConn, reader.Reader, metadata.Source.TCPAddr(), ws.StateServerSide)
+	conn = NewConn(wsConn, metadata.Source.TCPAddr(), ws.StateServerSide)
 	if len(earlyData) > 0 {
 		conn = bufio.NewCachedConn(conn, buf.As(earlyData))
 	}