From 1a1aa9d642864400f499a8f0d0e83a4b9546ba31 Mon Sep 17 00:00:00 2001
From: Joseph Doherty <dohejw01@gmail.com>
Date: Mon, 23 Feb 2026 04:47:57 -0500
Subject: [PATCH] fix: use byte-length for close message truncation, add
 exception-safe disposal

- CreateCloseMessage now operates on UTF-8 byte length (matching Go's
  len(body) behavior) instead of character length, with proper UTF-8
  boundary detection during truncation
- WsCompression.Compress now uses try/finally for exception-safe disposal
  of DeflateStream and MemoryStream
---
 src/NATS.Server/WebSocket/WsCompression.cs | 25 +++++++++++++---------
 src/NATS.Server/WebSocket/WsFrameWriter.cs | 25 ++++++++++++++++------
 2 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/src/NATS.Server/WebSocket/WsCompression.cs b/src/NATS.Server/WebSocket/WsCompression.cs
index d5b8a5d..92f0184 100644
--- a/src/NATS.Server/WebSocket/WsCompression.cs
+++ b/src/NATS.Server/WebSocket/WsCompression.cs
@@ -22,19 +22,24 @@ public static class WsCompression
     {
         var output = new MemoryStream();
         var deflate = new DeflateStream(output, CompressionLevel.Fastest, leaveOpen: true);
-        deflate.Write(data);
-        deflate.Flush();
+        try
+        {
+            deflate.Write(data);
+            deflate.Flush();
 
-        var compressed = output.ToArray();
+            var compressed = output.ToArray();
 
-        deflate.Dispose();
-        output.Dispose();
+            // Remove trailing 4-byte sync marker (0x00 0x00 0xff 0xff) per RFC 7692
+            if (compressed.Length >= 4)
+                return compressed[..^4];
 
-        // Remove trailing 4-byte sync marker (0x00 0x00 0xff 0xff) per RFC 7692
-        if (compressed.Length >= 4)
-            return compressed[..^4];
-
-        return compressed;
+            return compressed;
+        }
+        finally
+        {
+            deflate.Dispose();
+            output.Dispose();
+        }
     }
 
     /// <summary>
diff --git a/src/NATS.Server/WebSocket/WsFrameWriter.cs b/src/NATS.Server/WebSocket/WsFrameWriter.cs
index 1d0848d..59ba4f8 100644
--- a/src/NATS.Server/WebSocket/WsFrameWriter.cs
+++ b/src/NATS.Server/WebSocket/WsFrameWriter.cs
@@ -102,16 +102,27 @@ public static class WsFrameWriter
     /// </summary>
     public static byte[] CreateCloseMessage(int status, string body)
     {
-        if (body.Length > WsConstants.MaxControlPayloadSize - WsConstants.CloseStatusSize)
+        var bodyBytes = Encoding.UTF8.GetBytes(body);
+        int maxBody = WsConstants.MaxControlPayloadSize - WsConstants.CloseStatusSize;
+
+        if (bodyBytes.Length > maxBody)
         {
-            body = body[..(WsConstants.MaxControlPayloadSize - WsConstants.CloseStatusSize - 3)] + "...";
+            var suffix = "..."u8;
+            int truncLen = maxBody - suffix.Length;
+            // Find a valid UTF-8 boundary by walking back from truncation point
+            while (truncLen > 0 && (bodyBytes[truncLen] & 0xC0) == 0x80)
+                truncLen--;
+            var buf = new byte[WsConstants.CloseStatusSize + truncLen + suffix.Length];
+            BinaryPrimitives.WriteUInt16BigEndian(buf, (ushort)status);
+            bodyBytes.AsSpan(0, truncLen).CopyTo(buf.AsSpan(WsConstants.CloseStatusSize));
+            suffix.CopyTo(buf.AsSpan(WsConstants.CloseStatusSize + truncLen));
+            return buf;
         }
 
-        var bodyBytes = Encoding.UTF8.GetBytes(body);
-        var buf = new byte[WsConstants.CloseStatusSize + bodyBytes.Length];
-        BinaryPrimitives.WriteUInt16BigEndian(buf, (ushort)status);
-        bodyBytes.CopyTo(buf.AsSpan(WsConstants.CloseStatusSize));
-        return buf;
+        var result = new byte[WsConstants.CloseStatusSize + bodyBytes.Length];
+        BinaryPrimitives.WriteUInt16BigEndian(result, (ushort)status);
+        bodyBytes.CopyTo(result.AsSpan(WsConstants.CloseStatusSize));
+        return result;
     }
 
     /// <summary>