From 437b7e70e47027585f11a38abaa4e8dc7639d383 Mon Sep 17 00:00:00 2001
From: Ivan Zlatanov <ivan@zlatanov.net>
Date: Tue, 20 Apr 2021 11:18:40 +0300
Subject: [PATCH] Disallowing the usage of different compression options for
 continuations.

---
 .../System.Net.WebSockets/src/Resources/Strings.resx  |  3 +++
 .../src/System/Net/WebSockets/ManagedWebSocket.cs     |  8 +++++---
 .../System.Net.WebSockets/tests/WebSocketTests.cs     | 11 +++++++++++
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/src/libraries/System.Net.WebSockets/src/Resources/Strings.resx b/src/libraries/System.Net.WebSockets/src/Resources/Strings.resx
index c838d4d188ad45..693f8d3863fd7f 100644
--- a/src/libraries/System.Net.WebSockets/src/Resources/Strings.resx
+++ b/src/libraries/System.Net.WebSockets/src/Resources/Strings.resx
@@ -162,4 +162,7 @@
   <data name="ZLibUnsupportedCompression" xml:space="preserve">
     <value>The message was compressed using an unsupported compression method.</value>
   </data>
+  <data name="net_WebSockets_Argument_MessageFlagsHasDifferentCompressionOptions" xml:space="preserve">
+    <value>The compression options for a continuation cannot be different than the options used to send the first fragment of the message.</value>
+  </data>
 </root>
\ No newline at end of file
diff --git a/src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/ManagedWebSocket.cs b/src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/ManagedWebSocket.cs
index 01d633db8ca791..971c2ceff82be4 100644
--- a/src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/ManagedWebSocket.cs
+++ b/src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/ManagedWebSocket.cs
@@ -300,18 +300,20 @@ private ValueTask SendPrivateAsync(ReadOnlyMemory<byte> buffer, WebSocketMessage
             }
 
             bool endOfMessage = messageFlags.HasFlag(WebSocketMessageFlags.EndOfMessage);
-            bool disableCompression;
+            bool disableCompression = messageFlags.HasFlag(WebSocketMessageFlags.DisableCompression);
             MessageOpcode opcode;
 
             if (_lastSendWasFragment)
             {
-                disableCompression = _lastSendHadDisableCompression;
+                if (_lastSendHadDisableCompression != disableCompression)
+                {
+                    throw new ArgumentException(SR.net_WebSockets_Argument_MessageFlagsHasDifferentCompressionOptions, nameof(messageFlags));
+                }
                 opcode = MessageOpcode.Continuation;
             }
             else
             {
                 opcode = messageType == WebSocketMessageType.Binary ? MessageOpcode.Binary : MessageOpcode.Text;
-                disableCompression = messageFlags.HasFlag(WebSocketMessageFlags.DisableCompression);
             }
 
             ValueTask t = SendFrameAsync(opcode, endOfMessage, disableCompression, buffer, cancellationToken);
diff --git a/src/libraries/System.Net.WebSockets/tests/WebSocketTests.cs b/src/libraries/System.Net.WebSockets/tests/WebSocketTests.cs
index ad738a00ec8645..19b38d8d21760e 100644
--- a/src/libraries/System.Net.WebSockets/tests/WebSocketTests.cs
+++ b/src/libraries/System.Net.WebSockets/tests/WebSocketTests.cs
@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.IO;
+using System.Threading.Tasks;
 using Xunit;
 
 namespace System.Net.WebSockets.Tests
@@ -171,6 +172,16 @@ public void ValueWebSocketReceiveResult_Ctor_ValidArguments_Roundtrip(int count,
             Assert.Equal(endOfMessage, r.EndOfMessage);
         }
 
+        [Fact]
+        public async Task ThrowWhenContinuationWithDifferentCompressionFlags()
+        {
+            using WebSocket client = CreateFromStream(new MemoryStream(), isServer: false, null, TimeSpan.Zero);
+
+            await client.SendAsync(Memory<byte>.Empty, WebSocketMessageType.Text, WebSocketMessageFlags.DisableCompression, default);
+            Assert.Throws<ArgumentException>("messageFlags", () =>
+               client.SendAsync(Memory<byte>.Empty, WebSocketMessageType.Binary, WebSocketMessageFlags.EndOfMessage, default));
+        }
+
         public abstract class ExposeProtectedWebSocket : WebSocket
         {
             public static new bool IsStateTerminal(WebSocketState state) =>