Apply max_fragmented_message_size to the first fragment of a message

lib/bandit/websocket/connection.ex

@@ -110,8 +110,7 @@ defmodule Bandit.WebSocket.Connection do
         end
 
       %Frame.Text{fin: false} = frame ->
-        {:continue,
-         %{connection | fragment_frame: frame, fragment_size: IO.iodata_length(frame.data)}}
+        start_fragmented_message(frame, socket, connection)
 
       %Frame.Binary{fin: true, compressed: true} = frame ->
         do_inflate(frame, socket, connection)
@@ -121,8 +120,7 @@ defmodule Bandit.WebSocket.Connection do
         |> handle_continuation(socket, connection)
 
       %Frame.Binary{fin: false} = frame ->
-        {:continue,
-         %{connection | fragment_frame: frame, fragment_size: IO.iodata_length(frame.data)}}
+        start_fragmented_message(frame, socket, connection)
 
       frame ->
         handle_control_frame(frame, socket, connection)
@@ -186,6 +184,19 @@ defmodule Bandit.WebSocket.Connection do
     end
   end
 
+  # The first fragment of a message counts towards max_fragmented_message_size like the ones that
+  # follow it, so one larger than the limit is rejected as it arrives rather than being held until
+  # the next frame arrives and the limit is checked
+  defp start_fragmented_message(frame, socket, connection) do
+    fragment_size = IO.iodata_length(frame.data)
+
+    if oversize_message?(fragment_size, connection.opts) do
+      do_error(1009, "Received oversize fragmented message", socket, connection)
+    else
+      {:continue, %{connection | fragment_frame: frame, fragment_size: fragment_size}}
+    end
+  end
+
   defp maybe_coalesce_fragments(%{fragment_pending_count: count} = connection)
        when count >= @max_pending_fragments do
     chunk = IO.iodata_to_binary(connection.fragment_pending)

test/bandit/websocket/protocol_test.exs

@@ -354,6 +354,36 @@ defmodule WebSocketProtocolTest do
 
       assert output =~ "Received oversize fragmented message"
     end
+
+    # The first fragment is subject to the limit as well. It used to be stored unchecked, and the
+    # limit only applied once a continuation frame arrived
+    for {opcode, send_frame} <- [text: :send_text_frame, binary: :send_binary_frame] do
+      test "max_fragmented_message_size enforced for an opening #{opcode} fragment", context do
+        output =
+          capture_log(fn ->
+            context =
+              http_server(context, websocket_options: [max_fragmented_message_size: 1_000])
+
+            client = SimpleWebSocketClient.tcp_client(context)
+            SimpleWebSocketClient.http1_handshake(client, TerminateWebSock)
+
+            payload = String.duplicate("0123456789", 101)
+            apply(SimpleWebSocketClient, unquote(send_frame), [client, payload, 0x0])
+
+            # Get the error that terminate saw, to ensure we're closing for the expected reason
+            assert_receive {:error, "Received oversize fragmented message"}, 500
+
+            assert SimpleWebSocketClient.recv_connection_close_frame(client) ==
+                     {:ok, <<1009::16>>}
+
+            # Verify that the server didn't send any extraneous frames
+            assert SimpleWebSocketClient.connection_closed_for_reading?(client)
+            Process.sleep(500)
+          end)
+
+        assert output =~ "Received oversize fragmented message"
+      end
+    end
   end
 
   describe "compressed frames" do