From 4a0d097fe9f1b8fe352267040658331168e8abd8 Mon Sep 17 00:00:00 2001
From: wwqgtxx <wwqgtxx@gmail.com>
Date: Mon, 17 Apr 2023 00:23:12 +0800
Subject: [PATCH] fix: ensure StreamWebsocketConn call N.NewDeadlineConn

---
 common/net/bufconn.go        |  4 ++++
 transport/vmess/websocket.go | 14 +++++++++++---
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/common/net/bufconn.go b/common/net/bufconn.go
index 8087608c..f01a3dda 100644
--- a/common/net/bufconn.go
+++ b/common/net/bufconn.go
@@ -79,3 +79,7 @@ func (c *BufferedConn) ReaderReplaceable() bool {
 	}
 	return true
 }
+
+func (c *BufferedConn) WriterReplaceable() bool {
+	return true
+}
diff --git a/transport/vmess/websocket.go b/transport/vmess/websocket.go
index 179da886..e7335d84 100644
--- a/transport/vmess/websocket.go
+++ b/transport/vmess/websocket.go
@@ -334,7 +334,10 @@ func streamWebsocketWithEarlyDataConn(conn net.Conn, c *WebsocketConfig) (net.Co
 		underlay: conn,
 		config:   c,
 	}
-	return conn, nil
+	// websocketWithEarlyDataConn can't correct handle Deadline
+	// it will not apply the already set Deadline after Dial()
+	// so call N.NewDeadlineConn to add a safe wrapper
+	return N.NewDeadlineConn(conn), nil
 }
 
 func streamWebsocketConn(conn net.Conn, c *WebsocketConfig, earlyData *bytes.Buffer) (net.Conn, error) {
@@ -402,11 +405,16 @@ func streamWebsocketConn(conn net.Conn, c *WebsocketConfig, earlyData *bytes.Buf
 		return nil, fmt.Errorf("dial %s error: %s", uri.Host, reason)
 	}
 
-	return &websocketConn{
+	conn = &websocketConn{
 		conn:       wsConn,
 		rawWriter:  N.NewExtendedWriter(wsConn.UnderlyingConn()),
 		remoteAddr: conn.RemoteAddr(),
-	}, nil
+	}
+	// websocketConn can't correct handle ReadDeadline
+	// gorilla/websocket will cache the os.ErrDeadlineExceeded from conn.Read()
+	// it will cause read fail and event panic in *websocket.Conn.NextReader()
+	// so call N.NewDeadlineConn to add a safe wrapper
+	return N.NewDeadlineConn(conn), nil
 }
 
 func StreamWebsocketConn(conn net.Conn, c *WebsocketConfig) (net.Conn, error) {