Skip to content

Commit d6b3c4e

Browse files
authored
Merge pull request #181 from Fiekers/master
fix the blocking when reading websocket packets exceeding 1024 bytes
2 parents 52e1cba + b8e6d17 commit d6b3c4e

File tree

1 file changed

+21
-9
lines changed

1 file changed

+21
-9
lines changed

server/server.go

+21-9
Original file line numberDiff line numberDiff line change
@@ -1027,26 +1027,38 @@ var defaultUpgrader = &websocket.Upgrader{
10271027
Subprotocols: []string{"mqtt"},
10281028
}
10291029

1030-
//实现io.ReadWriter接口
1030+
// 实现io.ReadWriter接口
10311031
// wsConn implements the io.readWriter
10321032
type wsConn struct {
10331033
net.Conn
1034-
c *websocket.Conn
1034+
c *websocket.Conn
1035+
buf []byte
1036+
r int // buf copy positions
10351037
}
10361038

10371039
func (ws *wsConn) Close() error {
10381040
return ws.Conn.Close()
10391041
}
10401042

10411043
func (ws *wsConn) Read(p []byte) (n int, err error) {
1042-
msgType, r, err := ws.c.NextReader()
1043-
if err != nil {
1044-
return 0, err
1044+
if ws.buf == nil {
1045+
msgType, buf, err := ws.c.ReadMessage()
1046+
if err != nil {
1047+
return 0, err
1048+
}
1049+
if msgType != websocket.BinaryMessage {
1050+
return 0, ErrInvalWsMsgType
1051+
}
1052+
ws.buf = buf
10451053
}
1046-
if msgType != websocket.BinaryMessage {
1047-
return 0, ErrInvalWsMsgType
1054+
n = copy(p, ws.buf[ws.r:])
1055+
ws.r += n
1056+
// reset reader buffer
1057+
if ws.r+1 >= len(ws.buf) {
1058+
ws.buf = nil
1059+
ws.r = 0
10481060
}
1049-
return r.Read(p)
1061+
return
10501062
}
10511063

10521064
func (ws *wsConn) Write(p []byte) (n int, err error) {
@@ -1363,7 +1375,7 @@ func (srv *server) wsHandler() http.HandlerFunc {
13631375
return
13641376
}
13651377
defer c.Close()
1366-
conn := &wsConn{c.UnderlyingConn(), c}
1378+
conn := &wsConn{Conn: c.UnderlyingConn(), c: c}
13671379
client, err := srv.newClient(conn)
13681380
if err != nil {
13691381
zaplog.Error("new client fail", zap.Error(err))

0 commit comments

Comments
 (0)