Browse Source

Add the {deflate, boolean()} Websocket command

It allows to temporarily disable Websocket compression
when it was negotiated. It's ignored otherwise. This
can be used as fine-grained control when some frames
do not compress well.
Loïc Hoguin 6 years ago
parent
commit
8d920f3db9

+ 31 - 21
src/cowboy_websocket.erl

@@ -31,7 +31,10 @@
 -export([system_terminate/4]).
 -export([system_terminate/4]).
 -export([system_code_change/4]).
 -export([system_code_change/4]).
 
 
--type commands() :: [cow_ws:frame() | {active, boolean()}].
+-type commands() :: [cow_ws:frame()
+	| {active, boolean()}
+	| {deflate, boolean()}
+].
 -export_type([commands/0]).
 -export_type([commands/0]).
 
 
 -type call_result(State) :: {commands(), State} | {commands(), State, hibernate}.
 -type call_result(State) :: {commands(), State} | {commands(), State, hibernate}.
@@ -88,6 +91,7 @@
 	frag_state = undefined :: cow_ws:frag_state(),
 	frag_state = undefined :: cow_ws:frag_state(),
 	frag_buffer = <<>> :: binary(),
 	frag_buffer = <<>> :: binary(),
 	utf8_state = 0 :: cow_ws:utf8_state(),
 	utf8_state = 0 :: cow_ws:utf8_state(),
+	deflate = true :: boolean(),
 	extensions = #{} :: map(),
 	extensions = #{} :: map(),
 	req = #{} :: map()
 	req = #{} :: map()
 }).
 }).
@@ -424,10 +428,8 @@ parse_payload(State=#state{frag_state=FragState, utf8_state=Incomplete, extensio
 			websocket_close(State, HandlerState, Error)
 			websocket_close(State, HandlerState, Error)
 	end.
 	end.
 
 
-dispatch_frame(State=#state{opts=Opts, frag_state=FragState,
-		frag_buffer=SoFar, extensions=Extensions}, HandlerState,
-		#ps_payload{type=Type0, unmasked=Payload0, close_code=CloseCode0},
-		RemainingData) ->
+dispatch_frame(State=#state{opts=Opts, frag_state=FragState, frag_buffer=SoFar}, HandlerState,
+		#ps_payload{type=Type0, unmasked=Payload0, close_code=CloseCode0}, RemainingData) ->
 	MaxFrameSize = maps:get(max_frame_size, Opts, infinity),
 	MaxFrameSize = maps:get(max_frame_size, Opts, infinity),
 	case cow_ws:make_frame(Type0, Payload0, CloseCode0, FragState) of
 	case cow_ws:make_frame(Type0, Payload0, CloseCode0, FragState) of
 		%% @todo Allow receiving fragments.
 		%% @todo Allow receiving fragments.
@@ -446,12 +448,12 @@ dispatch_frame(State=#state{opts=Opts, frag_state=FragState,
 		{close, CloseCode, Payload} ->
 		{close, CloseCode, Payload} ->
 			websocket_close(State, HandlerState, {remote, CloseCode, Payload});
 			websocket_close(State, HandlerState, {remote, CloseCode, Payload});
 		Frame = ping ->
 		Frame = ping ->
-			transport_send(State, nofin, cow_ws:frame(pong, Extensions)),
+			transport_send(State, nofin, frame(pong, State)),
 			handler_call(State, HandlerState,
 			handler_call(State, HandlerState,
 				#ps_header{buffer=RemainingData},
 				#ps_header{buffer=RemainingData},
 				websocket_handle, Frame, fun parse_header/3);
 				websocket_handle, Frame, fun parse_header/3);
 		Frame = {ping, Payload} ->
 		Frame = {ping, Payload} ->
-			transport_send(State, nofin, cow_ws:frame({pong, Payload}, Extensions)),
+			transport_send(State, nofin, frame({pong, Payload}, State)),
 			handler_call(State, HandlerState,
 			handler_call(State, HandlerState,
 				#ps_header{buffer=RemainingData},
 				#ps_header{buffer=RemainingData},
 				websocket_handle, Frame, fun parse_header/3);
 				websocket_handle, Frame, fun parse_header/3);
@@ -523,8 +525,10 @@ commands([], State, Data) ->
 	{Result, State};
 	{Result, State};
 commands([{active, Active}|Tail], State, Data) when is_boolean(Active) ->
 commands([{active, Active}|Tail], State, Data) when is_boolean(Active) ->
 	commands(Tail, State#state{active=Active}, Data);
 	commands(Tail, State#state{active=Active}, Data);
-commands([Frame|Tail], State=#state{extensions=Extensions}, Data0) ->
-	Data = [cow_ws:frame(Frame, Extensions)|Data0],
+commands([{deflate, Deflate}|Tail], State, Data) when is_boolean(Deflate) ->
+	commands(Tail, State#state{deflate=Deflate}, Data);
+commands([Frame|Tail], State, Data0) ->
+	Data = [frame(Frame, State)|Data0],
 	case is_close_frame(Frame) of
 	case is_close_frame(Frame) of
 		true ->
 		true ->
 			_ = transport_send(State, fin, lists:reverse(Data)),
 			_ = transport_send(State, fin, lists:reverse(Data)),
@@ -542,8 +546,8 @@ transport_send(#state{socket=Socket, transport=Transport}, _, Data) ->
 -spec websocket_send(cow_ws:frame(), #state{}) -> ok | stop | {error, atom()}.
 -spec websocket_send(cow_ws:frame(), #state{}) -> ok | stop | {error, atom()}.
 websocket_send(Frames, State) when is_list(Frames) ->
 websocket_send(Frames, State) when is_list(Frames) ->
 	websocket_send_many(Frames, State, []);
 	websocket_send_many(Frames, State, []);
-websocket_send(Frame, State=#state{extensions=Extensions}) ->
-	Data = cow_ws:frame(Frame, Extensions),
+websocket_send(Frame, State) ->
+	Data = frame(Frame, State),
 	case is_close_frame(Frame) of
 	case is_close_frame(Frame) of
 		true ->
 		true ->
 			_ = transport_send(State, fin, Data),
 			_ = transport_send(State, fin, Data),
@@ -554,8 +558,8 @@ websocket_send(Frame, State=#state{extensions=Extensions}) ->
 
 
 websocket_send_many([], State, Acc) ->
 websocket_send_many([], State, Acc) ->
 	transport_send(State, nofin, lists:reverse(Acc));
 	transport_send(State, nofin, lists:reverse(Acc));
-websocket_send_many([Frame|Tail], State=#state{extensions=Extensions}, Acc0) ->
-	Acc = [cow_ws:frame(Frame, Extensions)|Acc0],
+websocket_send_many([Frame|Tail], State, Acc0) ->
+	Acc = [frame(Frame, State)|Acc0],
 	case is_close_frame(Frame) of
 	case is_close_frame(Frame) of
 		true ->
 		true ->
 			_ = transport_send(State, fin, lists:reverse(Acc)),
 			_ = transport_send(State, fin, lists:reverse(Acc)),
@@ -574,25 +578,31 @@ websocket_close(State, HandlerState, Reason) ->
 	websocket_send_close(State, Reason),
 	websocket_send_close(State, Reason),
 	terminate(State, HandlerState, Reason).
 	terminate(State, HandlerState, Reason).
 
 
-websocket_send_close(State=#state{extensions=Extensions}, Reason) ->
+websocket_send_close(State, Reason) ->
 	_ = case Reason of
 	_ = case Reason of
 		Normal when Normal =:= stop; Normal =:= timeout ->
 		Normal when Normal =:= stop; Normal =:= timeout ->
-			transport_send(State, fin, cow_ws:frame({close, 1000, <<>>}, Extensions));
+			transport_send(State, fin, frame({close, 1000, <<>>}, State));
 		{error, badframe} ->
 		{error, badframe} ->
-			transport_send(State, fin, cow_ws:frame({close, 1002, <<>>}, Extensions));
+			transport_send(State, fin, frame({close, 1002, <<>>}, State));
 		{error, badencoding} ->
 		{error, badencoding} ->
-			transport_send(State, fin, cow_ws:frame({close, 1007, <<>>}, Extensions));
+			transport_send(State, fin, frame({close, 1007, <<>>}, State));
 		{error, badsize} ->
 		{error, badsize} ->
-			transport_send(State, fin, cow_ws:frame({close, 1009, <<>>}, Extensions));
+			transport_send(State, fin, frame({close, 1009, <<>>}, State));
 		{crash, _, _} ->
 		{crash, _, _} ->
-			transport_send(State, fin, cow_ws:frame({close, 1011, <<>>}, Extensions));
+			transport_send(State, fin, frame({close, 1011, <<>>}, State));
 		remote ->
 		remote ->
-			transport_send(State, fin, cow_ws:frame(close, Extensions));
+			transport_send(State, fin, frame(close, State));
 		{remote, Code, _} ->
 		{remote, Code, _} ->
-			transport_send(State, fin, cow_ws:frame({close, Code, <<>>}, Extensions))
+			transport_send(State, fin, frame({close, Code, <<>>}, State))
 	end,
 	end,
 	ok.
 	ok.
 
 
+%% Don't compress frames while deflate is disabled.
+frame(Frame, #state{deflate=false, extensions=Extensions}) ->
+	cow_ws:frame(Frame, Extensions#{deflate => false});
+frame(Frame, #state{extensions=Extensions}) ->
+	cow_ws:frame(Frame, Extensions).
+
 -spec terminate(#state{}, any(), terminate_reason()) -> no_return().
 -spec terminate(#state{}, any(), terminate_reason()) -> no_return().
 terminate(State, HandlerState, Reason) ->
 terminate(State, HandlerState, Reason) ->
 	handler_terminate(State, HandlerState, Reason),
 	handler_terminate(State, HandlerState, Reason),

+ 2 - 2
test/handlers/ws_active_commands_h.erl

@@ -1,5 +1,5 @@
-%% This module takes commands from the x-commands header
-%% and returns them in the websocket_init/1 callback.
+%% This module starts with active mode disabled
+%% and enables it again once a timeout is triggered.
 
 
 -module(ws_active_commands_h).
 -module(ws_active_commands_h).
 -behavior(cowboy_websocket).
 -behavior(cowboy_websocket).

+ 24 - 0
test/handlers/ws_deflate_commands_h.erl

@@ -0,0 +1,24 @@
+%% This module enables/disables compression
+%% every time it echoes a frame.
+
+-module(ws_deflate_commands_h).
+-behavior(cowboy_websocket).
+
+-export([init/2]).
+-export([websocket_handle/2]).
+-export([websocket_info/2]).
+
+init(Req, RunOrHibernate) ->
+	{cowboy_websocket, Req,
+		#{deflate => true, hibernate => RunOrHibernate},
+		#{compress => true}}.
+
+websocket_handle(Frame, State=#{deflate := Deflate0, hibernate := run}) ->
+	Deflate = not Deflate0,
+	{[Frame, {deflate, Deflate}], State#{deflate => Deflate}};
+websocket_handle(Frame, State=#{deflate := Deflate0, hibernate := hibernate}) ->
+	Deflate = not Deflate0,
+	{[Frame, {deflate, Deflate}], State#{deflate => Deflate}, hibernate}.
+
+websocket_info(_Info, State) ->
+	{[], State}.

+ 41 - 1
test/ws_handler_SUITE.erl

@@ -50,7 +50,8 @@ init_dispatch(Name) ->
 		{"/init", ws_init_commands_h, RunOrHibernate},
 		{"/init", ws_init_commands_h, RunOrHibernate},
 		{"/handle", ws_handle_commands_h, RunOrHibernate},
 		{"/handle", ws_handle_commands_h, RunOrHibernate},
 		{"/info", ws_info_commands_h, RunOrHibernate},
 		{"/info", ws_info_commands_h, RunOrHibernate},
-		{"/active", ws_active_commands_h, RunOrHibernate}
+		{"/active", ws_active_commands_h, RunOrHibernate},
+		{"/deflate", ws_deflate_commands_h, RunOrHibernate}
 	]}]).
 	]}]).
 
 
 %% Support functions for testing using Gun.
 %% Support functions for testing using Gun.
@@ -217,3 +218,42 @@ websocket_active_false(Config) ->
 	{ok, {text, <<"Not received until the handler enables active again.">>}}
 	{ok, {text, <<"Not received until the handler enables active again.">>}}
 		= receive_ws(ConnPid, StreamRef),
 		= receive_ws(ConnPid, StreamRef),
 	ok.
 	ok.
+
+websocket_deflate_false(Config) ->
+	doc("The {deflate, false} command temporarily disables compression. "
+		"The {deflate, true} command reenables it."),
+	%% We disable context takeover so that the compressed data
+	%% does not change across all frames.
+	{ok, Socket, Headers} = ws_SUITE:do_handshake("/deflate",
+		"Sec-WebSocket-Extensions: permessage-deflate; server_no_context_takeover\r\n", Config),
+	{_, "permessage-deflate; server_no_context_takeover"}
+		= lists:keyfind("sec-websocket-extensions", 1, Headers),
+	%% The handler receives a compressed "Hello" frame and
+	%% sends back a compressed or uncompressed echo intermittently.
+	Mask = 16#11223344,
+	CompressedHello = <<242, 72, 205, 201, 201, 7, 0>>,
+	MaskedHello = ws_SUITE:do_mask(CompressedHello, Mask, <<>>),
+	%% First echo is compressed.
+	ok = gen_tcp:send(Socket, <<1:1, 1:1, 0:2, 1:4, 1:1, 7:7, Mask:32, MaskedHello/binary>>),
+	{ok, <<1:1, 1:1, 0:2, 1:4, 0:1, 7:7, CompressedHello/binary>>} = gen_tcp:recv(Socket, 0, 6000),
+	%% Second echo is not compressed when it is received back.
+	ok = gen_tcp:send(Socket, <<1:1, 1:1, 0:2, 1:4, 1:1, 7:7, Mask:32, MaskedHello/binary>>),
+	{ok, <<1:1, 0:3, 1:4, 0:1, 5:7, "Hello">>} = gen_tcp:recv(Socket, 0, 6000),
+	%% Third echo is compressed again.
+	ok = gen_tcp:send(Socket, <<1:1, 1:1, 0:2, 1:4, 1:1, 7:7, Mask:32, MaskedHello/binary>>),
+	{ok, <<1:1, 1:1, 0:2, 1:4, 0:1, 7:7, CompressedHello/binary>>} = gen_tcp:recv(Socket, 0, 6000),
+	%% Client-initiated close.
+	ok = gen_tcp:send(Socket, << 1:1, 0:3, 8:4, 1:1, 0:7, 0:32 >>),
+	{ok, << 1:1, 0:3, 8:4, 0:8 >>} = gen_tcp:recv(Socket, 0, 6000),
+	{error, closed} = gen_tcp:recv(Socket, 0, 6000),
+	ok.
+
+websocket_deflate_ignore_if_not_negotiated(Config) ->
+	doc("The {deflate, boolean()} commands are ignored "
+		"when compression was not negotiated."),
+	{ok, ConnPid, StreamRef} = gun_open_ws(Config, "/deflate", []),
+	_ = [begin
+		gun:ws_send(ConnPid, {text, <<"Hello.">>}),
+		{ok, {text, <<"Hello.">>}} = receive_ws(ConnPid, StreamRef)
+	end || _ <- lists:seq(1, 10)],
+	ok.