From 1bac18aeb1ab0a822ceac683b2cf5fd75a2435d5 Mon Sep 17 00:00:00 2001
From: Stef Heyenrath <stef.heyenrath@gmail.com>
Date: Wed, 11 Feb 2026 22:39:41 +0100
Subject: [PATCH] byte[]

---
 .github/copilot-instructions.md               |  4 ++++
 .../WebSockets/WebSocketMessageBuilder.cs     |  6 ++++++
 .../Extensions/ArrayPoolExtensions.cs         | 20 +++++++++++++++++++
 .../WebSockets/IWebSocketMessageBuilder.cs    |  6 ++++++
 .../WebSockets/ClientWebSocketExtensions.cs   |  7 ++++---
 5 files changed, 40 insertions(+), 3 deletions(-)
 create mode 100644 .github/copilot-instructions.md
 create mode 100644 src/WireMock.Net.Shared/Extensions/ArrayPoolExtensions.cs

diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md
new file mode 100644
index 00000000..aaf5be35
--- /dev/null
+++ b/.github/copilot-instructions.md
@@ -0,0 +1,4 @@
+# Copilot Instructions
+
+## Project Guidelines
+- All new byte[xx] calls should use using var data = ArrayPool<byte>.Shared.Lease(xx); instead of directly allocating byte arrays
\ No newline at end of file
diff --git a/src/WireMock.Net.Minimal/WebSockets/WebSocketMessageBuilder.cs b/src/WireMock.Net.Minimal/WebSockets/WebSocketMessageBuilder.cs
index 730908cc..81487690 100644
--- a/src/WireMock.Net.Minimal/WebSockets/WebSocketMessageBuilder.cs
+++ b/src/WireMock.Net.Minimal/WebSockets/WebSocketMessageBuilder.cs
@@ -58,6 +58,12 @@ internal class WebSocketMessageBuilder : IWebSocketMessageBuilder
         return this;
     }
 
+    public IWebSocketMessageBuilder Close()
+    {
+        ShouldClose = true;
+        return this;
+    }
+
     internal enum MessageType
     {
         Text,
diff --git a/src/WireMock.Net.Shared/Extensions/ArrayPoolExtensions.cs b/src/WireMock.Net.Shared/Extensions/ArrayPoolExtensions.cs
new file mode 100644
index 00000000..83ebcbcf
--- /dev/null
+++ b/src/WireMock.Net.Shared/Extensions/ArrayPoolExtensions.cs
@@ -0,0 +1,20 @@
+// Copyright © WireMock.Net
+
+namespace System.Buffers;
+
+internal sealed class Lease<T>(ArrayPool<T> pool, int length) : IDisposable
+{
+    public T[] Rented { get; } = pool.Rent(length);
+
+    public static implicit operator T[](Lease<T> lease) => lease.Rented;
+
+    public void Dispose()
+    {
+        pool.Return(Rented, true);
+    }
+}
+
+internal static class ArrayPoolExtensions
+{
+    public static Lease<T> Lease<T>(this ArrayPool<T> source, int length) => new(source, length);
+}
\ No newline at end of file
diff --git a/src/WireMock.Net.Shared/WebSockets/IWebSocketMessageBuilder.cs b/src/WireMock.Net.Shared/WebSockets/IWebSocketMessageBuilder.cs
index ee0b6fa5..118fd3cb 100644
--- a/src/WireMock.Net.Shared/WebSockets/IWebSocketMessageBuilder.cs
+++ b/src/WireMock.Net.Shared/WebSockets/IWebSocketMessageBuilder.cs
@@ -49,4 +49,10 @@ public interface IWebSocketMessageBuilder
     /// </summary>
     [PublicAPI]
     IWebSocketMessageBuilder AndClose();
+
+    /// <summary>
+    /// Close the WebSocket connection.
+    /// </summary>
+    [PublicAPI]
+    IWebSocketMessageBuilder Close();
 }
\ No newline at end of file
diff --git a/test/WireMock.Net.Tests/WebSockets/ClientWebSocketExtensions.cs b/test/WireMock.Net.Tests/WebSockets/ClientWebSocketExtensions.cs
index b7b4cac3..a0049b9b 100644
--- a/test/WireMock.Net.Tests/WebSockets/ClientWebSocketExtensions.cs
+++ b/test/WireMock.Net.Tests/WebSockets/ClientWebSocketExtensions.cs
@@ -1,5 +1,6 @@
 // Copyright © WireMock.Net
 
+using System.Buffers;
 using System.Net.WebSockets;
 using System.Text;
 
@@ -12,9 +13,9 @@ internal static class ClientWebSocketExtensions
         return client.SendAsync(new ArraySegment<byte>(Encoding.UTF8.GetBytes(text)), WebSocketMessageType.Text, endOfMessage, cancellationToken);
     }
 
-    internal static async Task<string> ReceiveAsTextAsync(this ClientWebSocket client, CancellationToken cancellationToken = default)
+    internal static async Task<string> ReceiveAsTextAsync(this ClientWebSocket client, int bufferSize = 1024, CancellationToken cancellationToken = default)
     {
-        var receiveBuffer = new byte[1024];
+        using var receiveBuffer = ArrayPool<byte>.Shared.Lease(1024);
         var result = await client.ReceiveAsync(new ArraySegment<byte>(receiveBuffer), cancellationToken);
 
         if (result.MessageType != WebSocketMessageType.Text)
@@ -32,7 +33,7 @@ internal static class ClientWebSocketExtensions
 
     internal static async Task<byte[]> ReceiveAsBytesAsync(this ClientWebSocket client, int bufferSize = 1024, CancellationToken cancellationToken = default)
     {
-        var receiveBuffer = new byte[bufferSize];
+        using var receiveBuffer = ArrayPool<byte>.Shared.Lease(1024);
         var result = await client.ReceiveAsync(new ArraySegment<byte>(receiveBuffer), cancellationToken);
 
         if (result.MessageType != WebSocketMessageType.Binary)