From 4c34774b7eb787e37892399f2daddba68ec277e3 Mon Sep 17 00:00:00 2001
From: Kirill Kinduk <evilbluebeaver@gmail.com>
Date: Fri, 25 Aug 2017 12:08:26 +0300
Subject: 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.
---
 src/cowboy_websocket.erl | 22 +++++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

(limited to 'src/cowboy_websocket.erl')

diff --git a/src/cowboy_websocket.erl b/src/cowboy_websocket.erl
index 725d7ec..df2e1a5 100644
--- a/src/cowboy_websocket.erl
+++ b/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,
-		ParseState=#ps_header{buffer=Data}) ->
+parse_header(State=#state{max_frame_size=MaxFrameSize,
+		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},
-		HandlerState, #ps_payload{type=Type0, unmasked=Payload0, close_code=CloseCode0},
+		max_frame_size=MaxFrameSize, frag_state=FragState,
+		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 ->
-- 
cgit v1.2.3