test gorilla websocket

transport/internet/websocket/connection.go

@@ -1,124 +1,91 @@
 package websocket
 
 import (
+	"context"
 	"io"
 	"net"
 	"time"
 
-	"github.com/gorilla/websocket"
+	"github.com/coder/websocket"
 	"github.com/xtls/xray-core/common/buf"
-	"github.com/xtls/xray-core/common/errors"
-	"github.com/xtls/xray-core/common/serial"
 )
 
 var _ buf.Writer = (*connection)(nil)
 
-// connection is a wrapper for net.Conn over WebSocket connection.
-// remoteAddr is used to pass "virtual" remote IP addresses in X-Forwarded-For.
-// so we shouldn't directly read it form conn.
 type connection struct {
-	conn       *websocket.Conn
+	net.Conn
 	reader     io.Reader
 	remoteAddr net.Addr
+	localAddr  net.Addr
 }
 
 func NewConnection(conn *websocket.Conn, remoteAddr net.Addr, extraReader io.Reader, heartbeatPeriod uint32) *connection {
+	ctx, cancel := context.WithCancel(context.Background())
+	c := &connection{Conn: websocket.NetConn(ctx, conn, websocket.MessageBinary), reader: extraReader, remoteAddr: remoteAddr}
+	if remoteAddr == nil {
+		c.remoteAddr = c.Conn.RemoteAddr()
+	}
+	c.Conn = &cancelConn{Conn: c.Conn, cancel: cancel}
 	if heartbeatPeriod != 0 {
 		go func() {
+			ticker := time.NewTicker(time.Duration(heartbeatPeriod) * time.Second)
+			defer ticker.Stop()
 			for {
-				time.Sleep(time.Duration(heartbeatPeriod) * time.Second)
-				if err := conn.WriteControl(websocket.PingMessage, []byte{}, time.Time{}); err != nil {
-					break
+				select {
+				case <-ctx.Done():
+					return
+				case <-ticker.C:
+					pingCtx, stop := context.WithTimeout(ctx, time.Duration(heartbeatPeriod)*time.Second)
+					err := conn.Ping(pingCtx)
+					stop()
+					if err != nil {
+						c.Close()
+						return
+					}
 				}
 			}
 		}()
 	}
-
-	return &connection{
-		conn:       conn,
-		remoteAddr: remoteAddr,
-		reader:     extraReader,
-	}
+	return c
+}
+
+type cancelConn struct {
+	net.Conn
+	cancel context.CancelFunc
+}
+
+func (c *cancelConn) Close() error {
+	defer c.cancel()
+	return c.Conn.Close()
 }
 
-// Read implements net.Conn.Read()
 func (c *connection) Read(b []byte) (int, error) {
-	for {
-		reader, err := c.getReader()
-		if err != nil {
-			return 0, err
-		}
-
-		nBytes, err := reader.Read(b)
-		if errors.Cause(err) == io.EOF {
-			c.reader = nil
-			continue
-		}
-		return nBytes, err
-	}
-}
-
-func (c *connection) getReader() (io.Reader, error) {
 	if c.reader != nil {
-		return c.reader, nil
+		n, err := c.reader.Read(b)
+		if err == io.EOF {
+			c.reader = nil
+			if n > 0 {
+				return n, nil
+			}
+		} else {
+			return n, err
+		}
 	}
-
-	_, reader, err := c.conn.NextReader()
-	if err != nil {
-		return nil, err
-	}
-	c.reader = reader
-	return reader, nil
+	return c.Conn.Read(b)
 }
 
-// Write implements io.Writer.
-func (c *connection) Write(b []byte) (int, error) {
-	if err := c.conn.WriteMessage(websocket.BinaryMessage, b); err != nil {
-		return 0, err
+func (c *connection) LocalAddr() net.Addr {
+	if c.localAddr != nil {
+		return c.localAddr
 	}
-	return len(b), nil
+	return c.Conn.LocalAddr()
 }
 
+func (c *connection) RemoteAddr() net.Addr { return c.remoteAddr }
+
 func (c *connection) WriteMultiBuffer(mb buf.MultiBuffer) error {
 	mb = buf.Compact(mb)
 	mb, err := buf.WriteMultiBuffer(c, mb)
 	buf.ReleaseMulti(mb)
 	return err
 }
-
-func (c *connection) Close() error {
-	var errs []interface{}
-	if err := c.conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second*5)); err != nil {
-		errs = append(errs, err)
-	}
-	if err := c.conn.Close(); err != nil {
-		errs = append(errs, err)
-	}
-	if len(errs) > 0 {
-		return errors.New("failed to close connection").Base(errors.New(serial.Concat(errs...)))
-	}
-	return nil
-}
-
-func (c *connection) LocalAddr() net.Addr {
-	return c.conn.LocalAddr()
-}
-
-func (c *connection) RemoteAddr() net.Addr {
-	return c.remoteAddr
-}
-
-func (c *connection) SetDeadline(t time.Time) error {
-	if err := c.SetReadDeadline(t); err != nil {
-		return err
-	}
-	return c.SetWriteDeadline(t)
-}
-
-func (c *connection) SetReadDeadline(t time.Time) error {
-	return c.conn.SetReadDeadline(t)
-}
-
-func (c *connection) SetWriteDeadline(t time.Time) error {
-	return c.conn.SetWriteDeadline(t)
-}

transport/internet/websocket/dialer.go

@@ -5,9 +5,10 @@ import (
 	_ "embed"
 	"encoding/base64"
 	"io"
+	"net/http"
 	"time"
 
-	"github.com/gorilla/websocket"
+	"github.com/coder/websocket"
 	"github.com/xtls/xray-core/common"
 	"github.com/xtls/xray-core/common/errors"
 	"github.com/xtls/xray-core/common/net"
@@ -46,23 +47,22 @@ func init() {
 func dialWebSocket(ctx context.Context, dest net.Destination, streamSettings *internet.MemoryStreamConfig, ed []byte) (net.Conn, error) {
 	wsSettings := streamSettings.ProtocolSettings.(*Config)
 
-	dialer := &websocket.Dialer{
-		NetDial: func(network, addr string) (net.Conn, error) {
+	var localAddr, remoteAddr net.Addr
+	dialer := &http.Transport{
+		DialContext: func(dialCtx context.Context, network, addr string) (net.Conn, error) {
 			var conn net.Conn
 			var err error
 			if streamSettings.FinalMask != nil {
-				conn, err = streamSettings.FinalMask.DialTCP(ctx, dest)
+				conn, err = streamSettings.FinalMask.DialTCP(dialCtx, dest)
 			} else {
-				conn, err = internet.DialSystem(ctx, dest, streamSettings.SocketSettings)
+				conn, err = internet.DialSystem(dialCtx, dest, streamSettings.SocketSettings)
 			}
 			if err != nil {
 				return nil, errors.New("failed to dial to dest").Base(err)
 			}
+			localAddr, remoteAddr = conn.LocalAddr(), conn.RemoteAddr()
 			return conn, err
 		},
-		ReadBufferSize:   4 * 1024,
-		WriteBufferSize:  4 * 1024,
-		HandshakeTimeout: time.Second * 8,
 	}
 
 	protocol := "ws"
@@ -73,14 +73,14 @@ func dialWebSocket(ctx context.Context, dest net.Destination, streamSettings *in
 		tlsConfig := tConfig.GetTLSConfig(tls.WithDestination(dest), tls.WithNextProto("http/1.1"))
 		dialer.TLSClientConfig = tlsConfig
 		if fingerprint := tls.GetFingerprint(tConfig.Fingerprint); fingerprint != nil {
-			dialer.NetDialTLSContext = func(_ context.Context, _, addr string) (net.Conn, error) {
+			dialer.DialTLSContext = func(dialCtx context.Context, _, addr string) (net.Conn, error) {
 				// Like the NetDial in the dialer
 				var pconn net.Conn
 				var err error
 				if streamSettings.FinalMask != nil {
-					pconn, err = streamSettings.FinalMask.DialTCP(ctx, dest)
+					pconn, err = streamSettings.FinalMask.DialTCP(dialCtx, dest)
 				} else {
-					pconn, err = internet.DialSystem(ctx, dest, streamSettings.SocketSettings)
+					pconn, err = internet.DialSystem(dialCtx, dest, streamSettings.SocketSettings)
 				}
 				if err != nil {
 					return nil, errors.New("failed to dial to dest").Base(err)
@@ -88,16 +88,19 @@ func dialWebSocket(ctx context.Context, dest net.Destination, streamSettings *in
 
 				// TLS and apply the handshake
 				cn := tls.UClient(pconn, tlsConfig, fingerprint).(*tls.UConn)
-				if err := cn.WebsocketHandshakeContext(ctx); err != nil {
+				if err := cn.WebsocketHandshakeContext(dialCtx); err != nil {
+					pconn.Close()
 					errors.LogErrorInner(ctx, err, "failed to dial to "+addr)
 					return nil, err
 				}
 				if !tlsConfig.InsecureSkipVerify {
 					if err := cn.VerifyHostname(tlsConfig.ServerName); err != nil {
+						pconn.Close()
 						errors.LogErrorInner(ctx, err, "failed to dial to "+addr)
 						return nil, err
 					}
 				}
+				localAddr, remoteAddr = cn.LocalAddr(), cn.RemoteAddr()
 				return cn, nil
 			}
 		}
@@ -106,7 +109,7 @@ func dialWebSocket(ctx context.Context, dest net.Destination, streamSettings *in
 	if browser_dialer.HasBrowserDialer() {
 		// For Browser Dialer's optimized IP and non-standard port
 		host := wsSettings.Host
-		if host == "" && tConfig.ServerName != "" {
+		if host == "" && tConfig != nil && tConfig.ServerName != "" {
 			host = tConfig.ServerName
 		}
 		if host == "" {
@@ -122,7 +125,7 @@ func dialWebSocket(ctx context.Context, dest net.Destination, streamSettings *in
 			return nil, err
 		}
 
-		return NewConnection(conn, conn.RemoteAddr(), nil, wsSettings.HeartbeatPeriod), nil
+		return NewConnection(conn, nil, nil, wsSettings.HeartbeatPeriod), nil
 	}
 
 	host := dest.Address.String()
@@ -145,7 +148,18 @@ func dialWebSocket(ctx context.Context, dest net.Destination, streamSettings *in
 		header.Set("Sec-WebSocket-Protocol", base64.RawURLEncoding.EncodeToString(ed))
 	}
 
-	conn, resp, err := dialer.DialContext(ctx, uri, header)
+	subprotocols := []string(nil)
+	if protocol := header.Get("Sec-WebSocket-Protocol"); protocol != "" {
+		subprotocols = []string{protocol}
+		header.Del("Sec-WebSocket-Protocol")
+	}
+	hostHeader := header.Get("Host")
+	header.Del("Host")
+	dialCtx, cancel := context.WithTimeout(ctx, 8*time.Second)
+	defer cancel()
+	conn, resp, err := websocket.Dial(dialCtx, uri, &websocket.DialOptions{
+		HTTPClient: &http.Client{Transport: dialer}, HTTPHeader: header, Host: hostHeader, Subprotocols: subprotocols,
+	})
 	if err != nil {
 		var reason string
 		if resp != nil {
@@ -154,7 +168,9 @@ func dialWebSocket(ctx context.Context, dest net.Destination, streamSettings *in
 		return nil, errors.New("failed to dial to (", uri, "): ", reason).Base(err)
 	}
 
-	return NewConnection(conn, conn.RemoteAddr(), nil, wsSettings.HeartbeatPeriod), nil
+	wrapped := NewConnection(conn, remoteAddr, nil, wsSettings.HeartbeatPeriod)
+	wrapped.localAddr = localAddr
+	return wrapped, nil
 }
 
 type delayDialConn struct {

transport/internet/websocket/hub.go

@@ -11,7 +11,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/gorilla/websocket"
+	"github.com/coder/websocket"
 	"github.com/xtls/xray-core/common"
 	"github.com/xtls/xray-core/common/errors"
 	"github.com/xtls/xray-core/common/net"
@@ -29,15 +29,6 @@ type requestHandler struct {
 
 var replacer = strings.NewReplacer("+", "-", "/", "_", "=", "")
 
-var upgrader = &websocket.Upgrader{
-	ReadBufferSize:   0,
-	WriteBufferSize:  0,
-	HandshakeTimeout: time.Second * 4,
-	CheckOrigin: func(r *http.Request) bool {
-		return true
-	},
-}
-
 func (h *requestHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
 	if len(h.host) > 0 && !internet.IsValidHTTPHost(request.Host, h.host) {
 		errors.LogInfo(context.Background(), "failed to validate host, request:", request.Host, ", config:", h.host)
@@ -51,28 +42,28 @@ func (h *requestHandler) ServeHTTP(writer http.ResponseWriter, request *http.Req
 	}
 
 	var extraReader io.Reader
-	responseHeader := http.Header{}
+	acceptOptions := &websocket.AcceptOptions{InsecureSkipVerify: true}
 	if str := request.Header.Get("Sec-WebSocket-Protocol"); str != "" {
 		if ed, err := base64.RawURLEncoding.DecodeString(replacer.Replace(str)); err == nil && len(ed) > 0 {
 			extraReader = bytes.NewReader(ed)
-			responseHeader.Set("Sec-WebSocket-Protocol", str)
+			acceptOptions.Subprotocols = []string{str}
 		}
 	}
 
-	conn, err := upgrader.Upgrade(writer, request, responseHeader)
+	conn, err := websocket.Accept(writer, request, acceptOptions)
 	if err != nil {
 		errors.LogInfoInner(context.Background(), err, "failed to convert to WebSocket connection")
 		return
 	}
 
-	remoteAddr := conn.RemoteAddr()
+	remoteAddr, _ := net.ResolveTCPAddr("tcp", request.RemoteAddr)
 	var trustedXFF []string
 	if h.socketSettings != nil {
 		trustedXFF = h.socketSettings.TrustedXForwardedFor
 	}
-	remoteAddr = http_proto.ApplyTrustedXForwardedFor(request.Header, trustedXFF, remoteAddr)
+	effectiveRemoteAddr := http_proto.ApplyTrustedXForwardedFor(request.Header, trustedXFF, remoteAddr)
 
-	h.ln.addConn(NewConnection(conn, remoteAddr, extraReader, h.ln.config.HeartbeatPeriod))
+	h.ln.addConn(NewConnection(conn, effectiveRemoteAddr, extraReader, h.ln.config.HeartbeatPeriod))
 }
 
 type Listener struct {