Add max_frame_size option for websocket handlers

Option allows to limit a frame by size before decoding its payload.

LH: I have added a test for when the limit is reached on a nofin
fragmented frame (the last commit addressed that case but it had
no test). I have fixed formatting and other, and changed the
default value to infinity since it might otherwise be incompatible
with existing code. I also added documentation and a bunch of other
minor changes.

doc/src/guide/ws_handlers.asciidoc

@@ -225,6 +225,22 @@ init(Req, State) ->
 This value cannot be changed once it is set. It defaults to
 `60000`.
 
+=== Limiting frame sizes
+
+Cowboy accepts frames of any size by default. You should
+limit the size depending on what your handler may handle.
+You can do this via the `init/2` callback:
+
+[source,erlang]
+----
+init(Req, State) ->
+    {cowboy_websocket, Req, State, #{
+        max_frame_size => 8000000}}.
+----
+
+The lack of limit is historical. A future version of
+Cowboy will have a more reasonable default.
+
 === Saving memory
 
 The Websocket connection process can be set to hibernate

doc/src/manual/cowboy_websocket.asciidoc

@@ -153,6 +153,7 @@ Cowboy does it automatically for you.
 opts() :: #{
     compress => boolean(),
     idle_timeout => timeout(),
+    max_frame_size => non_neg_integer() | infinity,
     req_filter => fun((cowboy_req:req()) -> map())
 }
 ----
@@ -181,6 +182,13 @@ idle_timeout (60000)::
     connection open without receiving anything from
     the client.
 
+max_frame_size (infinity)::
+    Maximum frame size allowed by this Websocket
+    handler. Cowboy will close the connection when
+    a client attempts to send a frame that goes over
+    this limit. For fragmented frames this applies
+    to the size of the reconstituted frame.
+
 req_filter::
     A function applied to the Req to compact it and
     only keep required information. The Req is only

src/cowboy_websocket.erl

@@ -57,6 +57,7 @@
 -type opts() :: #{
 	compress => boolean(),
 	idle_timeout => timeout(),
+	max_frame_size => non_neg_integer() | infinity,
 	req_filter => fun((cowboy_req:req()) -> map())
 }.
 -export_type([opts/0]).
@@ -71,6 +72,7 @@
 	timeout = infinity :: timeout(),
 	timeout_ref = undefined :: undefined | reference(),
 	compress = false :: boolean(),
+	max_frame_size :: non_neg_integer() | infinity,
 	messages = undefined :: undefined | {atom(), atom(), atom()},
 	hibernate = false :: boolean(),
 	frag_state = undefined :: cow_ws:frag_state(),
@@ -95,12 +97,14 @@ upgrade(Req, Env, Handler, HandlerState) ->
 %% @todo Error out if HTTP/2.
 upgrade(Req0, Env, Handler, HandlerState, Opts) ->
 	Timeout = maps:get(idle_timeout, Opts, 60000),
+	MaxFrameSize = maps:get(max_frame_size, Opts, infinity),
 	Compress = maps:get(compress, Opts, false),
 	FilteredReq = case maps:get(req_filter, Opts, undefined) of
 		undefined -> maps:with([method, version, scheme, host, port, path, qs, peer], Req0);
 		FilterFun -> FilterFun(Req0)
 	end,
-	State0 = #state{handler=Handler, timeout=Timeout, compress=Compress, req=FilteredReq},
+	State0 = #state{handler=Handler, timeout=Timeout, compress=Compress,
+		max_frame_size=MaxFrameSize, req=FilteredReq},
 	try websocket_upgrade(State0, Req0) of
 		{ok, State, Req} ->
 			websocket_handshake(State, Req, HandlerState, Env);
@@ -291,12 +295,15 @@ parse(State, HandlerState, PS=#ps_payload{buffer=Buffer}, Data) ->
 	parse_payload(State, HandlerState, PS#ps_payload{buffer= <<>>},
 		<<Buffer/binary, Data/binary>>).
 
-parse_header(State=#state{frag_state=FragState, extensions=Extensions}, HandlerState,
+parse_header(State=#state{max_frame_size=MaxFrameSize,
-		ParseState=#ps_header{buffer=Data}) ->
+		frag_state=FragState, extensions=Extensions},
+		HandlerState, ParseState=#ps_header{buffer=Data}) ->
 	case cow_ws:parse_header(Data, Extensions, FragState) of
 		%% All frames sent from the client to the server are masked.
 		{_, _, _, _, undefined, _} ->
 			websocket_close(State, HandlerState, {error, badframe});
+		{_, _, _, Len, _, _} when Len > MaxFrameSize ->
+			websocket_close(State, HandlerState, {error, badsize});
 		{Type, FragState2, Rsv, Len, MaskKey, Rest} ->
 			parse_payload(State#state{frag_state=FragState2}, HandlerState,
 				#ps_payload{type=Type, len=Len, mask_key=MaskKey, rsv=Rsv}, Rest);
@@ -335,11 +342,14 @@ parse_payload(State=#state{frag_state=FragState, utf8_state=Incomplete, extensio
 	end.
 
 dispatch_frame(State=#state{socket=Socket, transport=Transport,
-		frag_state=FragState, frag_buffer=SoFar, extensions=Extensions},
+		max_frame_size=MaxFrameSize, frag_state=FragState,
-		HandlerState, #ps_payload{type=Type0, unmasked=Payload0, close_code=CloseCode0},
+		frag_buffer=SoFar, extensions=Extensions}, HandlerState,
+		#ps_payload{type=Type0, unmasked=Payload0, close_code=CloseCode0},
 		RemainingData) ->
 	case cow_ws:make_frame(Type0, Payload0, CloseCode0, FragState) of
 		%% @todo Allow receiving fragments.
+		{fragment, _, _, Payload} when byte_size(Payload) + byte_size(SoFar) > MaxFrameSize ->
+			websocket_close(State, HandlerState, {error, badsize});
 		{fragment, nofin, _, Payload} ->
 			parse_header(State#state{frag_buffer= << SoFar/binary, Payload/binary >>},
 				HandlerState, #ps_header{buffer=RemainingData});
@@ -447,6 +457,8 @@ websocket_send_close(#state{socket=Socket, transport=Transport,
 			Transport:send(Socket, cow_ws:frame({close, 1002, <<>>}, Extensions));
 		{error, badencoding} ->
 			Transport:send(Socket, cow_ws:frame({close, 1007, <<>>}, Extensions));
+		{error, badsize} ->
+			Transport:send(Socket, cow_ws:frame({close, 1009, <<>>}, Extensions));
 		{crash, _, _} ->
 			Transport:send(Socket, cow_ws:frame({close, 1011, <<>>}, Extensions));
 		remote ->

test/ws_SUITE.erl

@@ -82,7 +82,8 @@ init_dispatch() ->
 			{"/ws_subprotocol", ws_subprotocol, []},
 			{"/terminate", ws_terminate_h, []},
 			{"/ws_timeout_hibernate", ws_timeout_hibernate, []},
-			{"/ws_timeout_cancel", ws_timeout_cancel, []}
+			{"/ws_timeout_cancel", ws_timeout_cancel, []},
+			{"/ws_max_frame_size", ws_max_frame_size, []}
 		]}
 	]).
 
@@ -302,6 +303,44 @@ ws_init_shutdown_before_handshake(Config) ->
 	{ok, {http_response, {1, 1}, 403, _}, _Rest} = erlang:decode_packet(http, Handshake, []),
 	ok.
 
+ws_max_frame_size_close(Config) ->
+	doc("Server closes connection when frame size exceeds max_frame_size option"),
+	%% max_frame_size is set to 8 bytes in ws_max_frame_size.
+	{ok, Socket, _} = do_handshake("/ws_max_frame_size", Config),
+	Mask = 16#11223344,
+	MaskedHello = do_mask(<<"HelloHello">>, Mask, <<>>),
+	ok = gen_tcp:send(Socket, << 1:1, 0:3, 2:4, 1:1, 10:7, Mask:32, MaskedHello/binary >>),
+	{ok, << 1:1, 0:3, 8:4, 0:1, 2:7, 1009:16 >>} = gen_tcp:recv(Socket, 0, 6000),
+	{error, closed} = gen_tcp:recv(Socket, 0, 6000),
+	ok.
+
+ws_max_frame_size_final_fragment_close(Config) ->
+	doc("Server closes connection when final fragmented frame "
+		"exceeds max_frame_size option"),
+	%% max_frame_size is set to 8 bytes in ws_max_frame_size.
+	{ok, Socket, _} = do_handshake("/ws_max_frame_size", Config),
+	Mask = 16#11223344,
+	MaskedHello = do_mask(<<"Hello">>, Mask, <<>>),
+	ok = gen_tcp:send(Socket, << 0:1, 0:3, 2:4, 1:1, 5:7, Mask:32, MaskedHello/binary >>),
+	ok = gen_tcp:send(Socket, << 1:1, 0:3, 0:4, 1:1, 5:7, Mask:32, MaskedHello/binary >>),
+	{ok, << 1:1, 0:3, 8:4, 0:1, 2:7, 1009:16 >>} = gen_tcp:recv(Socket, 0, 6000),
+	{error, closed} = gen_tcp:recv(Socket, 0, 6000),
+	ok.
+
+ws_max_frame_size_intermediate_fragment_close(Config) ->
+	doc("Server closes connection when intermediate fragmented frame "
+		"exceeds max_frame_size option"),
+	%% max_frame_size is set to 8 bytes in ws_max_frame_size.
+	{ok, Socket, _} = do_handshake("/ws_max_frame_size", Config),
+	Mask = 16#11223344,
+	MaskedHello = do_mask(<<"Hello">>, Mask, <<>>),
+	ok = gen_tcp:send(Socket, << 0:1, 0:3, 2:4, 1:1, 5:7, Mask:32, MaskedHello/binary >>),
+	ok = gen_tcp:send(Socket, << 0:1, 0:3, 0:4, 1:1, 5:7, Mask:32, MaskedHello/binary >>),
+	ok = gen_tcp:send(Socket, << 1:1, 0:3, 0:4, 1:1, 5:7, Mask:32, MaskedHello/binary >>),
+	{ok, << 1:1, 0:3, 8:4, 0:1, 2:7, 1009:16 >>} = gen_tcp:recv(Socket, 0, 6000),
+	{error, closed} = gen_tcp:recv(Socket, 0, 6000),
+	ok.
+
 ws_send_close(Config) ->
 	doc("Server-initiated close frame ends the connection."),
 	{ok, Socket, _} = do_handshake("/ws_send_close", Config),
@@ -402,8 +441,8 @@ ws_text_fragments(Config) ->
 	%% Send three "Hello" over three fragments and one send.
 	ok = gen_tcp:send(Socket, [
 		<< 0:1, 0:3, 1:4, 1:1, 5:7, Mask:32, MaskedHello/binary >>,
-		<< 0:1, 0:3, 0:4, 1:1, 5:7, Mask:32, MaskedHello/binary  >>,
+		<< 0:1, 0:3, 0:4, 1:1, 5:7, Mask:32, MaskedHello/binary >>,
-		<< 1:1, 0:3, 0:4, 1:1, 5:7, Mask:32, MaskedHello/binary  >>]),
+		<< 1:1, 0:3, 0:4, 1:1, 5:7, Mask:32, MaskedHello/binary >>]),
 	{ok, << 1:1, 0:3, 1:4, 0:1, 15:7, "HelloHelloHello" >>} = gen_tcp:recv(Socket, 0, 6000),
 	ok.
 

test/ws_SUITE_data/ws_max_frame_size.erl

@@ -0,0 +1,22 @@
+-module(ws_max_frame_size).
+
+-export([init/2]).
+-export([websocket_init/1]).
+-export([websocket_handle/2]).
+-export([websocket_info/2]).
+
+init(Req, State) ->
+	{cowboy_websocket, Req, State, #{max_frame_size => 8}}.
+
+websocket_init(State) ->
+	{ok, State}.
+
+websocket_handle({text, Data}, State) ->
+	{reply, {text, Data}, State};
+websocket_handle({binary, Data}, State) ->
+	{reply, {binary, Data}, State};
+websocket_handle(_Frame, State) ->
+	{ok, State}.
+
+websocket_info(_Info, State) ->
+	{ok, State}.