diff --git a/lib/bandit/websocket/connection.ex b/lib/bandit/websocket/connection.ex index cbfa563ed..f6c78cd9a 100644 --- a/lib/bandit/websocket/connection.ex +++ b/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) diff --git a/test/bandit/websocket/protocol_test.exs b/test/bandit/websocket/protocol_test.exs index b9cfd6c51..0b93cb67a 100644 --- a/test/bandit/websocket/protocol_test.exs +++ b/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