Add :max_message_size option to limit incoming message size

lib/mint/web_socket.ex

@@ -128,12 +128,16 @@ defmodule Mint.WebSocket do
             extensions: [Extension.t()],
             fragment: tuple(),
             private: map(),
-            buffer: binary()
+            buffer: binary(),
+            max_message_size: non_neg_integer() | :infinity
           }
   defstruct extensions: [],
             fragment: nil,
             private: %{},
-            buffer: <<>>
+            buffer: <<>>,
+            max_message_size: :infinity
+
+  @default_max_message_size 8_000_000
 
   @type error :: Mint.Types.error() | WebSocketError.t() | UpgradeFailureError.t()
 
@@ -322,6 +326,12 @@ defmodule Mint.WebSocket do
     * `:mode` - (default: `:active`) either `:active` or `:passive`. This
       corresponds to the same option in `Mint.HTTP.connect/4`.
 
+    * `:max_message_size` - (default: `#{@default_max_message_size}`) the
+      maximum size in bytes of an incoming message, after decompression.
+      `decode/2` returns a `Mint.WebSocketError` with reason
+      `:message_too_large` for bigger messages. Use `:infinity` to disable
+      the limit.
+
   ## Examples
 
       http_reply = receive(do: (message -> message))
@@ -342,13 +352,23 @@ defmodule Mint.WebSocket do
           {:ok, Mint.HTTP.t(), t()} | {:error, Mint.HTTP.t(), error()}
   def new(conn, request_ref, status, response_headers, opts \\ []) do
     websockets = [request_ref | get_private(conn, :websockets) || []]
+    max_message_size = Keyword.get(opts, :max_message_size, @default_max_message_size)
+
+    # extensions which buffer or inflate data (e.g. permessage-deflate) read the limit from their opts
+    extensions =
+      for extension <- get_private(conn, :extensions) do
+        %{extension | opts: Keyword.put(extension.opts, :max_message_size, max_message_size)}
+      end
 
     conn =
       conn
       |> put_private(:websockets, websockets)
       |> put_private(:mode, Keyword.get(opts, :mode, :active))
+      |> put_private(:extensions, extensions)
 
-    do_new(protocol(conn), conn, status, response_headers)
+    with {:ok, conn, websocket} <- do_new(protocol(conn), conn, status, response_headers) do
+      {:ok, conn, %{websocket | max_message_size: max_message_size}}
+    end
   end
 
   defp do_new(:http1, conn, status, headers) when status != 101 do

lib/mint/web_socket/frame.ex

@@ -242,7 +242,7 @@ defmodule Mint.WebSocket.Frame do
          websocket,
          acc
        ) do
-    case decode_payload_and_mask(payload_and_mask, masked == 0b1) do
+    case decode_payload_and_mask(payload_and_mask, masked == 0b1, websocket.max_message_size) do
       {:ok, payload, mask, rest} ->
         frame = decode_full_frame_binary(opcode, fin, reserved, mask, payload)
 
@@ -262,8 +262,9 @@ defmodule Mint.WebSocket.Frame do
     {:buffer, partial, :lists.reverse(acc)}
   end
 
-  defp decode_payload_and_mask(payload, masked?) do
+  defp decode_payload_and_mask(payload, masked?, max_message_size) do
     with {:ok, payload_length, rest} <- decode_payload_length(payload),
+         :ok <- check_message_size(payload_length, max_message_size),
          {:ok, mask, rest} <- decode_mask(rest, masked?),
          <<payload::binary-size(^payload_length), more::bitstring>> <- rest do
       {:ok, payload, mask, more}
@@ -274,6 +275,12 @@ defmodule Mint.WebSocket.Frame do
     end
   end
 
+  @doc false
+  def check_message_size(size, max) when is_integer(max) and size > max,
+    do: throw({:mint, %WebSocketError{reason: :message_too_large}})
+
+  def check_message_size(_size, _max), do: :ok
+
   defp decode_full_frame_binary(opcode, fin, reserved, mask, payload) do
     with {:ok, opcode} <- decode_opcode(opcode) do
       into_frame(
@@ -459,14 +466,14 @@ defmodule Mint.WebSocket.Frame do
   end
 
   def resolve_fragments(websocket, [frame | rest], acc) when is_fin(frame) do
-    frame = combine(websocket.fragment, frame)
+    frame = combine(websocket.fragment, frame, websocket.max_message_size)
 
     put_in(websocket.fragment, nil)
     |> resolve_fragments(rest, [frame | acc])
   end
 
   def resolve_fragments(websocket, [frame | rest], acc) do
-    case combine(websocket.fragment, frame) do
+    case combine(websocket.fragment, frame, websocket.max_message_size) do
       {:error, reason} ->
         put_in(websocket.fragment, nil)
         |> resolve_fragments(rest, [{:error, reason} | acc])
@@ -477,18 +484,23 @@ defmodule Mint.WebSocket.Frame do
     end
   end
 
-  defp combine(nil, continuation(fin?: true)), do: {:error, :insular_continuation}
+  defp combine(nil, continuation(fin?: true), _max_message_size),
+    do: {:error, :insular_continuation}
 
-  defp combine(nil, frame), do: frame
+  defp combine(nil, frame, _max_message_size), do: frame
 
   for type <- [:continuation, :text, :binary] do
     defp combine(
            unquote(type)(data: frame_data) = frame,
-           continuation(data: continuation_data, fin?: fin?)
+           continuation(data: continuation_data, fin?: fin?),
+           max_message_size
          ) do
+      :ok =
+        check_message_size(byte_size(frame_data) + byte_size(continuation_data), max_message_size)
+
       unquote(type)(frame, data: Utils.maybe_concat(frame_data, continuation_data), fin?: fin?)
     end
   end
 
-  defp combine(a, b), do: {:error, {:cannot_combine_frames, a, b}}
+  defp combine(a, b, _max_message_size), do: {:error, {:cannot_combine_frames, a, b}}
 end

lib/mint/web_socket/per_message_deflate.ex

@@ -24,10 +24,11 @@ defmodule Mint.WebSocket.PerMessageDeflate do
           inflate: :zlib.zstream(),
           deflate: :zlib.zstream(),
           inflate_takeover?: boolean(),
-          deflate_takeover?: boolean()
+          deflate_takeover?: boolean(),
+          max_message_size: non_neg_integer() | :infinity
         }
 
-  defstruct [:inflate, :deflate, :inflate_takeover?, :deflate_takeover?]
+  defstruct [:inflate, :deflate, :inflate_takeover?, :deflate_takeover?, :max_message_size]
 
   @behaviour Extension
 
@@ -59,7 +60,8 @@ defmodule Mint.WebSocket.PerMessageDeflate do
       inflate: inflate_zstream,
       deflate: deflate_zstream,
       inflate_takeover?: get_takeover(params, "server_no_context_takeover", true),
-      deflate_takeover?: get_takeover(params, "client_no_context_takeover", true)
+      deflate_takeover?: get_takeover(params, "client_no_context_takeover", true),
+      max_message_size: Keyword.get(opts, :max_message_size, :infinity)
     }
 
     {:ok, put_in(this_extension.state, state)}
@@ -84,10 +86,7 @@ defmodule Mint.WebSocket.PerMessageDeflate do
       # Append 4 octets of 0x00 0x00 0xff 0xff to the tail end of the
       # payload of the message
 
-      data =
-        state.inflate
-        |> :zlib.inflate(<<data::binary, 0x00, 0x00, 0xFF, 0xFF>>)
-        |> IO.iodata_to_binary()
+      data = inflate(state, <<data::binary, 0x00, 0x00, 0xFF, 0xFF>>)
 
       if state.inflate_takeover? == false do
         :zlib.inflateReset(state.inflate)
@@ -105,6 +104,19 @@ defmodule Mint.WebSocket.PerMessageDeflate do
 
   def decode(frame, state), do: {:ok, frame, state}
 
+  # Known limitation: the limit is checked after each chunk returned by
+  # :zlib.safeInflate/2, so up to one chunk (16 KB on OTP 27) may be inflated past it.
+  defp inflate(state, data, acc \\ [], size \\ 0) do
+    {status, output} = :zlib.safeInflate(state.inflate, data)
+    size = size + IO.iodata_length(output)
+    :ok = Frame.check_message_size(size, state.max_message_size)
+
+    case status do
+      :continue -> inflate(state, [], [acc | output], size)
+      :finished -> IO.iodata_to_binary([acc | output])
+    end
+  end
+
   @doc false
   @impl Extension
   def encode(frame, state)

lib/mint/web_socket_error.ex

@@ -10,6 +10,7 @@ defmodule Mint.WebSocketError do
     quote do
       :extended_connect_disabled
       | :payload_too_large
+      | :message_too_large
       | {:extension_not_negotiated, Mint.WebSocket.Extension.t()}
     end
 
@@ -30,6 +31,10 @@ defmodule Mint.WebSocketError do
     "frame payload cannot exceed 9,223,372,036,854,775,807 bytes"
   end
 
+  defp format_reason(:message_too_large) do
+    "message exceeds the :max_message_size option"
+  end
+
   defp format_reason({:extension_not_negotiated, extension}) do
     "the remote server accepted an extension the client did not offer: #{inspect(extension)}"
   end

test/fixtures/autobahn_client.ex

@@ -82,7 +82,9 @@ defmodule AutobahnClient do
         [{:done, ^ref}] -> <<>>
       end
 
-    {:ok, conn, websocket} = Mint.WebSocket.new(conn, ref, status, resp_headers)
+    # the performance cases send messages up to 16 MiB
+    {:ok, conn, websocket} =
+      Mint.WebSocket.new(conn, ref, status, resp_headers, max_message_size: :infinity)
 
     %__MODULE__{
       next: :cont,

test/mint/web_socket_test.exs

@@ -56,6 +56,98 @@ defmodule Mint.WebSocketTest do
     end
   end
 
+  describe "given an HTTP/1 connection with :max_message_size" do
+    setup do
+      host = System.get_env("ECHO_HOST") || "localhost"
+      {:ok, conn} = HTTP1.connect(:http, host, 9000)
+      {:ok, conn, ref} = WebSocket.upgrade(:ws, conn, "/", [])
+      assert_receive http_get_message
+
+      {:ok, conn, [{:status, ^ref, status}, {:headers, ^ref, resp_headers}, {:done, ^ref}]} =
+        WebSocket.stream(conn, http_get_message)
+
+      {:ok, conn, websocket} =
+        WebSocket.new(conn, ref, status, resp_headers, max_message_size: 100)
+
+      on_exit(fn -> HTTP1.close(conn) end)
+      [websocket: websocket]
+    end
+
+    test "a frame declaring a length over the limit is rejected from its header",
+         %{websocket: websocket} do
+      # text frame header declaring a 4 GB payload, followed by a few bytes
+      header = <<1::1, 0::3, 1::4, 0::1, 127::7, 4_294_967_296::64>>
+
+      assert {:error, _websocket, %Mint.WebSocketError{reason: :message_too_large}} =
+               WebSocket.decode(websocket, header <> "abc")
+    end
+
+    test "a fragmented message over the limit is rejected", %{websocket: websocket} do
+      # 60 bytes each: first text frame (FIN = 0), then the final continuation (FIN = 1)
+      first = <<0::1, 0::3, 1::4, 0::1, 60::7>> <> :binary.copy("a", 60)
+      last = <<1::1, 0::3, 0::4, 0::1, 60::7>> <> :binary.copy("a", 60)
+
+      assert {:ok, websocket, []} = WebSocket.decode(websocket, first)
+
+      assert {:error, _websocket, %Mint.WebSocketError{reason: :message_too_large}} =
+               WebSocket.decode(websocket, last)
+    end
+
+    test "a fragmented message of exactly the limit is delivered", %{websocket: websocket} do
+      first = <<0::1, 0::3, 1::4, 0::1, 60::7>> <> :binary.copy("a", 60)
+      last = <<1::1, 0::3, 0::4, 0::1, 40::7>> <> :binary.copy("a", 40)
+
+      assert {:ok, websocket, []} = WebSocket.decode(websocket, first)
+      assert {:ok, _websocket, [{:text, message}]} = WebSocket.decode(websocket, last)
+      assert byte_size(message) == 100
+    end
+  end
+
+  describe "given an HTTP/1 connection with compression and :max_message_size" do
+    setup do
+      {:ok, conn} = HTTP1.connect(:http, "localhost", 7070)
+
+      {:ok, conn, ref} =
+        WebSocket.upgrade(:ws, conn, "/", [], extensions: [WebSocket.PerMessageDeflate])
+
+      assert_receive http_get_message
+
+      {:ok, conn, [{:status, ^ref, status}, {:headers, ^ref, resp_headers}, {:done, ^ref}]} =
+        WebSocket.stream(conn, http_get_message)
+
+      {:ok, conn, websocket} =
+        WebSocket.new(conn, ref, status, resp_headers, max_message_size: 100)
+
+      on_exit(fn -> HTTP1.close(conn) end)
+      [conn: conn, ref: ref, websocket: websocket]
+    end
+
+    test "a compressed message inflating over the limit is rejected", c do
+      # 1,000 identical bytes compress to a frame far below the limit
+      {:ok, websocket, data} = WebSocket.encode(c.websocket, {:text, :binary.copy("a", 1_000)})
+      {:ok, conn} = WebSocket.stream_request_body(c.conn, c.ref, data)
+
+      assert_receive echo_message
+      {:ok, _conn, [{:data, _ref, data}]} = WebSocket.stream(conn, echo_message)
+
+      assert byte_size(data) < 100
+
+      assert {:error, _websocket, %Mint.WebSocketError{reason: :message_too_large}} =
+               WebSocket.decode(websocket, data)
+    end
+
+    test "a compressed message inflating exactly to the limit is delivered", c do
+      message = :binary.copy("a", 100)
+      {:ok, websocket, data} = WebSocket.encode(c.websocket, {:text, message})
+      {:ok, conn} = WebSocket.stream_request_body(c.conn, c.ref, data)
+
+      assert_receive echo_message
+      {:ok, _conn, [{:data, _ref, data}]} = WebSocket.stream(conn, echo_message)
+
+      assert {:ok, _websocket, [{:text, ^message}]} = WebSocket.decode(websocket, data)
+    end
+  end
+
   describe "given a passive HTTP/1 connection to an echo server" do
     setup do
       host = System.get_env("ECHO_HOST") || "localhost"