From 0b29deaf27313199a5036b4bac2abec1079690fd Mon Sep 17 00:00:00 2001
From: Neil Twigg <neil@nats.io>
Date: Fri, 8 Nov 2024 11:41:08 +0000
Subject: [PATCH] WebSocket: Ensure full writes to compressor

Signed-off-by: Neil Twigg <neil@nats.io>
---
 server/websocket.go | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/server/websocket.go b/server/websocket.go
index 8f601a5477..7ebbaa85b7 100644
--- a/server/websocket.go
+++ b/server/websocket.go
@@ -1367,7 +1367,19 @@ func (c *client) wsCollapsePtoNB() (net.Buffers, int64) {
 		}
 		var csz int
 		for _, b := range nb {
-			cp.Write(b)
+			for len(b) > 0 {
+				n, err := cp.Write(b)
+				if err != nil {
+					if err == io.EOF {
+						break
+					}
+					c.Errorf("Error during compression: %v", err)
+					c.markConnAsClosed(WriteError)
+					nbPoolPut(b)
+					return nil, 0
+				}
+				b = b[n:]
+			}
 			nbPoolPut(b) // No longer needed as contents written to compressor.
 		}
 		if err := cp.Flush(); err != nil {