Repository: mochi-mqtt/server (Go)

Issue:
Websocket listener falls into infinite read loop for messages > ~2048 bytes
The byte buffer passed to `Read` is 2048 bytes and the current websocket Listener implementation incorrectly handles the case where the total number of bytes read exceeds this buffer size.  Specifically, the for-loop continually calls `r.Read` with an empty byte slice, which of course returns `(0, nil)` every time -- leading to a continual loop of retries.


Discussion:

Patch:
diff --git a/listeners/websocket.go b/listeners/websocket.go
index 0a5df58d..4e1f4d82 100644
--- a/listeners/websocket.go
+++ b/listeners/websocket.go
@@ -99,7 +99,7 @@ func (l *Websocket) handler(w http.ResponseWriter, r *http.Request) {
 	}
 	defer c.Close()
 
-	err = l.establish(l.id, &wsConn{c.UnderlyingConn(), c})
+	err = l.establish(l.id, &wsConn{Conn: c.UnderlyingConn(), c: c})
 	if err != nil {
 		l.log.Warn().Err(err).Send()
 	}
@@ -135,25 +135,41 @@ func (l *Websocket) Close(closeClients CloseFn) {
 type wsConn struct {
 	net.Conn
 	c *websocket.Conn
+
+	// reader for the current message (may be nil)
+	r io.Reader
 }
 
 // Read reads the next span of bytes from the websocket connection and returns the number of bytes read.
 func (ws *wsConn) Read(p []byte) (int, error) {
-	op, r, err := ws.c.NextReader()
-	if err != nil {
-		return 0, err
-	}
+	if ws.r == nil {
+		op, r, err := ws.c.NextReader()
+		if err != nil {
+			return 0, err
+		}
 
-	if op != websocket.BinaryMessage {
-		err = ErrInvalidMessage
-		return 0, err
+		if op != websocket.BinaryMessage {
+			err = ErrInvalidMessage
+			return 0, err
+		}
+
+		ws.r = r
 	}
 
-	var n, br int
+	var n int
 	for {
-		br, err = r.Read(p[n:])
+		// buffer is full, return what we've read so far
+		if n == len(p) {
+			return n, nil
+		}
+
+		br, err := ws.r.Read(p[n:])
 		n += br
 		if err != nil {
+			// when ANY error occurs, we consider this the end of the current message (either because it really is, via
+			// io.EOF, or because something bad happened, in which case we want to drop the remainder)
+			ws.r = nil
+
 			if errors.Is(err, io.EOF) {
 				err = nil
 			}


Test patch:
diff --git a/listeners/websocket_test.go b/listeners/websocket_test.go
index 03d3548e..ee91f81b 100644
--- a/listeners/websocket_test.go
+++ b/listeners/websocket_test.go
@@ -112,3 +112,44 @@ func TestWebsocketUpgrade(t *testing.T) {
 	s.Close()
 	ws.Close()
 }
+
+func TestWebsocketConnectionReads(t *testing.T) {
+	l := NewWebsocket("t1", testAddr, nil)
+	l.Init(nil)
+
+	recv := make(chan []byte)
+	l.establish = func(id string, c net.Conn) error {
+		var out []byte
+		for {
+			buf := make([]byte, 2048)
+			n, err := c.Read(buf)
+			require.NoError(t, err)
+			out = append(out, buf[:n]...)
+			if n < 2048 {
+				break
+			}
+		}
+
+		recv <- out
+		return nil
+	}
+
+	s := httptest.NewServer(http.HandlerFunc(l.handler))
+	ws, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(s.URL, "http"), nil)
+	require.NoError(t, err)
+
+	pkt := make([]byte, 3000) // make sure this is >2048
+	for i := 0; i < len(pkt); i++ {
+		pkt[i] = byte(i % 100)
+	}
+
+	err = ws.WriteMessage(websocket.BinaryMessage, pkt)
+	require.NoError(t, err)
+
+	got := <-recv
+	require.Equal(t, 3000, len(got))
+	require.Equal(t, pkt, got)
+
+	s.Close()
+	ws.Close()
+}