From fc9ba139386a323992f5907e019cfe25a5a9b9fc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Hoguin?= <essen@ninenines.eu>
Date: Thu, 16 Jan 2025 14:40:37 +0100
Subject: [PATCH] Websocket: Allow setting the max_frame_size option
 dynamically

This can be used to limit the maximum frame size before
some authentication or other validation is completed.
---
 src/cowboy_websocket.erl                    | 16 +++++-----
 test/handlers/ws_set_options_commands_h.erl | 19 ++++++++---
 test/ws_handler_SUITE.erl                   | 35 +++++++++++++++++++++
 3 files changed, 59 insertions(+), 11 deletions(-)

diff --git a/src/cowboy_websocket.erl b/src/cowboy_websocket.erl
index 577de47d..12c99bad 100644
--- a/src/cowboy_websocket.erl
+++ b/src/cowboy_websocket.erl
@@ -615,14 +615,16 @@ commands([{active, Active}|Tail], State0=#state{active=Active0}, Data) when is_b
 	commands(Tail, State#state{active=Active}, Data);
 commands([{deflate, Deflate}|Tail], State, Data) when is_boolean(Deflate) ->
 	commands(Tail, State#state{deflate=Deflate}, Data);
-commands([{set_options, SetOpts}|Tail], State0=#state{opts=Opts}, Data) ->
-	State = case SetOpts of
-		#{idle_timeout := IdleTimeout} ->
+commands([{set_options, SetOpts}|Tail], State0, Data) ->
+	State = maps:fold(fun
+		(idle_timeout, IdleTimeout, StateF=#state{opts=Opts}) ->
 			%% We reset the number of ticks when changing the idle_timeout option.
-			set_idle_timeout(State0#state{opts=Opts#{idle_timeout => IdleTimeout}}, 0);
-		_ ->
-			State0
-	end,
+			set_idle_timeout(StateF#state{opts=Opts#{idle_timeout => IdleTimeout}}, 0);
+		(max_frame_size, MaxFrameSize, StateF=#state{opts=Opts}) ->
+			StateF#state{opts=Opts#{max_frame_size => MaxFrameSize}};
+		(_, _, StateF) ->
+			StateF
+	end, State0, SetOpts),
 	commands(Tail, State, Data);
 commands([{shutdown_reason, ShutdownReason}|Tail], State, Data) ->
 	commands(Tail, State#state{shutdown_reason=ShutdownReason}, Data);
diff --git a/test/handlers/ws_set_options_commands_h.erl b/test/handlers/ws_set_options_commands_h.erl
index 88d4e722..1ab0af4d 100644
--- a/test/handlers/ws_set_options_commands_h.erl
+++ b/test/handlers/ws_set_options_commands_h.erl
@@ -11,10 +11,21 @@ init(Req, RunOrHibernate) ->
 	{cowboy_websocket, Req, RunOrHibernate,
 		#{idle_timeout => infinity}}.
 
-websocket_handle(Frame={text, <<"idle_timeout_short">>}, State=run) ->
-	{[{set_options, #{idle_timeout => 500}}, Frame], State};
-websocket_handle(Frame={text, <<"idle_timeout_short">>}, State=hibernate) ->
-	{[{set_options, #{idle_timeout => 500}}, Frame], State, hibernate}.
+%% Set the idle_timeout option dynamically.
+websocket_handle({text, <<"idle_timeout_short">>}, State=run) ->
+	{[{set_options, #{idle_timeout => 500}}], State};
+websocket_handle({text, <<"idle_timeout_short">>}, State=hibernate) ->
+	{[{set_options, #{idle_timeout => 500}}], State, hibernate};
+%% Set the max_frame_size option dynamically.
+websocket_handle({text, <<"max_frame_size_small">>}, State=run) ->
+	{[{set_options, #{max_frame_size => 1000}}], State};
+websocket_handle({text, <<"max_frame_size_small">>}, State=hibernate) ->
+	{[{set_options, #{max_frame_size => 1000}}], State, hibernate};
+%% We just echo binary frames.
+websocket_handle(Frame={binary, _}, State=run) ->
+	{[Frame], State};
+websocket_handle(Frame={binary, _}, State=hibernate) ->
+	{[Frame], State, hibernate}.
 
 websocket_info(_Info, State) ->
 	{[], State}.
diff --git a/test/ws_handler_SUITE.erl b/test/ws_handler_SUITE.erl
index ab1ffc84..3b842977 100644
--- a/test/ws_handler_SUITE.erl
+++ b/test/ws_handler_SUITE.erl
@@ -296,6 +296,41 @@ websocket_set_options_idle_timeout(Config) ->
 		error(timeout)
 	end.
 
+websocket_set_options_max_frame_size(Config) ->
+	doc("The max_frame_size option can be modified using the "
+		"command {set_options, Opts} at runtime."),
+	ConnPid = gun_open(Config),
+	StreamRef = gun:ws_upgrade(ConnPid, "/set_options"),
+	receive
+		{gun_upgrade, ConnPid, StreamRef, [<<"websocket">>], _} ->
+			ok;
+		{gun_response, ConnPid, _, _, Status, Headers} ->
+			exit({ws_upgrade_failed, Status, Headers});
+		{gun_error, ConnPid, StreamRef, Reason} ->
+			exit({ws_upgrade_failed, Reason})
+	after 1000 ->
+		error(timeout)
+	end,
+	%% We first send a 1MB frame to confirm that yes, we can
+	%% send a frame that large. The default max_frame_size is infinity.
+	gun:ws_send(ConnPid, StreamRef, {binary, <<0:8000000>>}),
+	{ws, {binary, <<0:8000000>>}} = gun:await(ConnPid, StreamRef),
+	%% Trigger the change in max_frame_size. From now on we will
+	%% only allow frames of up to 1000 bytes.
+	gun:ws_send(ConnPid, StreamRef, {text, <<"max_frame_size_small">>}),
+	%% Confirm that we can send frames of up to 1000 bytes.
+	gun:ws_send(ConnPid, StreamRef, {binary, <<0:8000>>}),
+	{ws, {binary, <<0:8000>>}} = gun:await(ConnPid, StreamRef),
+	%% Confirm that sending frames larger than 1000 bytes
+	%% results in the closing of the connection.
+	gun:ws_send(ConnPid, StreamRef, {binary, <<0:8008>>}),
+	receive
+		{gun_down, ConnPid, _, _, _} ->
+			ok
+	after 2000 ->
+		error(timeout)
+	end.
+
 websocket_shutdown_reason(Config) ->
 	doc("The command {shutdown_reason, any()} can be used to "
 		"change the shutdown reason of a Websocket connection."),