diff --git a/lib/mint/web_socket.ex b/lib/mint/web_socket.ex index 01839afb..36b7fa29 100644 --- a/lib/mint/web_socket.ex +++ b/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 diff --git a/lib/mint/web_socket/frame.ex b/lib/mint/web_socket/frame.ex index 89243331..66bec5d8 100644 --- a/lib/mint/web_socket/frame.ex +++ b/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?), <> <- 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 diff --git a/lib/mint/web_socket/per_message_deflate.ex b/lib/mint/web_socket/per_message_deflate.ex index 069d7721..fb69caad 100644 --- a/lib/mint/web_socket/per_message_deflate.ex +++ b/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(<>) - |> IO.iodata_to_binary() + data = inflate(state, <>) 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) diff --git a/lib/mint/web_socket_error.ex b/lib/mint/web_socket_error.ex index 7d9c5d8b..855df163 100644 --- a/lib/mint/web_socket_error.ex +++ b/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 diff --git a/test/fixtures/autobahn_client.ex b/test/fixtures/autobahn_client.ex index 7876705c..4ee55153 100644 --- a/test/fixtures/autobahn_client.ex +++ b/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, diff --git a/test/mint/web_socket_test.exs b/test/mint/web_socket_test.exs index fdc9ed9f..19885aff 100644 --- a/test/mint/web_socket_test.exs +++ b/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"