[release/6.0] Fix compression (#79549)

* Fix compression

* Apply suggestions from code review

Co-authored-by: Miha Zupan <mihazupan.zupan1@gmail.com>

* Adding SendAsync to ref

* fix ws deflate tests

* Check bytes on server side

* Fix u8

Co-authored-by: Miha Zupan <mihazupan.zupan1@gmail.com>

* Fix build on net6.0

* Remove ref assembly change, pr feedback

Co-authored-by: Katya Sokolova <esokolov@microsoft.com>
Co-authored-by: Miha Zupan <mihazupan.zupan1@gmail.com>
Co-authored-by: Natalia Kondratyeva <knatalia@microsoft.com>

src/libraries/System.Net.WebSockets.Client/src/System/Net/WebSockets/ClientWebSocket.cs

@@ -109,6 +109,9 @@ public override Task SendAsync(ArraySegment<byte> buffer, WebSocketMessageType m
         public override ValueTask SendAsync(ReadOnlyMemory<byte> buffer, WebSocketMessageType messageType, bool endOfMessage, CancellationToken cancellationToken) =>
             ConnectedWebSocket.SendAsync(buffer, messageType, endOfMessage, cancellationToken);
 
+        public override ValueTask SendAsync(ReadOnlyMemory<byte> buffer, WebSocketMessageType messageType, WebSocketMessageFlags messageFlags, CancellationToken cancellationToken) =>
+            ConnectedWebSocket.SendAsync(buffer, messageType, messageFlags, cancellationToken);
+
         public override Task<WebSocketReceiveResult> ReceiveAsync(ArraySegment<byte> buffer, CancellationToken cancellationToken) =>
             ConnectedWebSocket.ReceiveAsync(buffer, cancellationToken);
 

src/libraries/System.Net.WebSockets.Client/tests/DeflateTests.cs

@@ -72,6 +72,100 @@ await LoopbackServer.CreateClientAndServerAsync(async uri =>
             }), new LoopbackServer.Options { WebSocketEndpoint = true });
         }
 
+        [ConditionalFact(nameof(WebSocketsSupported))]
+        public async Task ThrowsWhenContinuationHasDifferentCompressionFlags()
+        {
+            var deflateOpt = new WebSocketDeflateOptions
+            {
+                ClientMaxWindowBits = 14,
+                ClientContextTakeover = true,
+                ServerMaxWindowBits = 14,
+                ServerContextTakeover = true
+            };
+            await LoopbackServer.CreateClientAndServerAsync(async uri =>
+            {
+                using var cws = new ClientWebSocket();
+                using var cts = new CancellationTokenSource(TimeOutMilliseconds);
+
+                cws.Options.DangerousDeflateOptions = deflateOpt;
+                await cws.ConnectAsync(uri, cts.Token);
+
+                await cws.SendAsync(Memory<byte>.Empty, WebSocketMessageType.Text, WebSocketMessageFlags.DisableCompression, default);
+                Assert.Throws<ArgumentException>("messageFlags", () =>
+                   cws.SendAsync(Memory<byte>.Empty, WebSocketMessageType.Binary, WebSocketMessageFlags.EndOfMessage, default));
+            }, server => server.AcceptConnectionAsync(async connection =>
+            {
+                string extensionsReply = CreateDeflateOptionsHeader(deflateOpt);
+                await LoopbackHelper.WebSocketHandshakeAsync(connection, extensionsReply);
+            }), new LoopbackServer.Options { WebSocketEndpoint = true });
+        }
+
+        [ConditionalFact(nameof(WebSocketsSupported))]
+        public async Task SendHelloWithDisableCompression()
+        {
+            byte[] bytes = Encoding.ASCII.GetBytes("Hello");
+
+            int prefixLength = 2;
+            byte[] rawPrefix = new byte[] { 0x81, 0x85 }; // fin=1, rsv=0, opcode=text; mask=1, len=5
+            int rawRemainingBytes = 9; // mask bytes (4) + payload bytes (5)
+            byte[] compressedPrefix = new byte[] { 0xc1, 0x87 }; // fin=1, rsv=compressed, opcode=text; mask=1, len=7
+            int compressedRemainingBytes = 11; // mask bytes (4) + payload bytes (7)
+
+            var deflateOpt = new WebSocketDeflateOptions
+            {
+                ClientMaxWindowBits = 14,
+                ClientContextTakeover = true,
+                ServerMaxWindowBits = 14,
+                ServerContextTakeover = true
+            };
+
+            await LoopbackServer.CreateClientAndServerAsync(async uri =>
+            {
+                using var cws = new ClientWebSocket();
+                using var cts = new CancellationTokenSource(TimeOutMilliseconds);
+
+                cws.Options.DangerousDeflateOptions = deflateOpt;
+                await cws.ConnectAsync(uri, cts.Token);
+
+                await cws.SendAsync(bytes, WebSocketMessageType.Text, true, cts.Token);
+
+                WebSocketMessageFlags flags = WebSocketMessageFlags.DisableCompression | WebSocketMessageFlags.EndOfMessage;
+                await cws.SendAsync(bytes, WebSocketMessageType.Text, flags, cts.Token);
+            }, server => server.AcceptConnectionAsync(async connection =>
+            {
+                var buffer = new byte[compressedRemainingBytes];
+                string extensionsReply = CreateDeflateOptionsHeader(deflateOpt);
+                await LoopbackHelper.WebSocketHandshakeAsync(connection, extensionsReply);
+
+                // first message is compressed
+                await ReadExactAsync(buffer, prefixLength);
+                Assert.Equal(compressedPrefix, buffer[..prefixLength]);
+                // read rest of the frame
+                await ReadExactAsync(buffer, compressedRemainingBytes);
+
+                // second message is not compressed
+                await ReadExactAsync(buffer, prefixLength);
+                Assert.Equal(rawPrefix, buffer[..prefixLength]);
+                // read rest of the frame
+                await ReadExactAsync(buffer, rawRemainingBytes);
+
+                async Task ReadExactAsync(byte[] buf, int n)
+                {
+                    var mem = buf.AsMemory(0, n);
+                    int totalRead = 0;
+                    while (totalRead < n)
+                    {
+                        int read = await connection.Stream.ReadAsync(mem.Slice(totalRead)).ConfigureAwait(false);
+                        if (read == 0)
+                        {
+                            throw new Exception("Unexpected end of stream");
+                        }
+                        totalRead += read;
+                    }
+                }
+            }), new LoopbackServer.Options { WebSocketEndpoint = true });
+        }
+
         private static string CreateDeflateOptionsHeader(WebSocketDeflateOptions options)
         {
             var builder = new StringBuilder();

src/libraries/System.Net.WebSockets.Client/tests/wasm/BrowserTimerThrottlingTest.cs

@@ -158,12 +158,12 @@ private async static Task SendAndReceive(ClientWebSocket cws, string message)
             catch (OperationCanceledException)
             {
             }
+#if DEBUG
             catch (Exception ex)
             {
-#if DEBUG
                 Console.WriteLine("SendAndReceive fail:" + ex);
-#endif
             }
+#endif
         }
     }