Disallowing the usage of different compression options for continuati…

…ons.
dotnet · Apr 26, 2021 · 437b7e7 · 437b7e7
1 parent 935969a
commit 437b7e7
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 3 deletions.
diff --git 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>
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
@@ -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
@@ -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) =>