Reject requests/data when using Websocket

3 files changed, 38 insertions, 8 deletions

diff --git a/src/gun.erl b/src/gun.erl
index ee3a5f6..ab26dbf 100644
--- a/src/gun.erl
+++ b/src/gun.erl
@@ -98,9 +98,10 @@
 -export([connecting/3]).
 -export([initial_tls_handshake/3]).
 -export([tls_handshake/3]).
--export([connected_no_input/3]).
--export([connected_data_only/3]).
 -export([connected/3]).
+-export([connected_data_only/3]).
+-export([connected_no_input/3]).
+-export([connected_ws_only/3]).
 -export([closing/3]).
 -export([terminate/3]).
 
@@ -1068,9 +1069,33 @@ protocol_negotiated({error, protocol_not_negotiated}, _) -> http.
 connected_no_input(Type, Event, State) ->
 	handle_common_connected_no_input(Type, Event, ?FUNCTION_NAME, State).
 
+connected_data_only(cast, Msg, _)
+		when element(1, Msg) =:= headers; element(1, Msg) =:= request;
+			element(1, Msg) =:= connect; element(1, Msg) =:= ws_upgrade;
+			element(1, Msg) =:= ws_send ->
+	ReplyTo = element(2, Msg),
+	ReplyTo ! {gun_error, self(), {badstate,
+		"This connection does not accept new requests to be opened "
+		"nor does it accept Websocket frames."}},
+	keep_state_and_data;
 connected_data_only(Type, Event, State) ->
 	handle_common_connected(Type, Event, ?FUNCTION_NAME, State).
 
+connected_ws_only(cast, {ws_send, Owner, Frames}, State=#state{
+		owner=Owner, protocol=Protocol=gun_ws, protocol_state=ProtoState,
+		event_handler=EvHandler, event_handler_state=EvHandlerState0}) ->
+	{Commands, EvHandlerState} = Protocol:send(Frames, ProtoState, EvHandler, EvHandlerState0),
+	commands(Commands, State#state{event_handler_state=EvHandlerState});
+connected_ws_only(cast, Msg, _)
+		when element(1, Msg) =:= headers; element(1, Msg) =:= request; element(1, Msg) =:= data;
+			element(1, Msg) =:= connect; element(1, Msg) =:= ws_upgrade ->
+	ReplyTo = element(2, Msg),
+	ReplyTo ! {gun_error, self(), {badstate,
+		"This connection only accepts Websocket frames."}},
+	keep_state_and_data;
+connected_ws_only(Type, Event, State) ->
+	handle_common_connected_no_input(Type, Event, ?FUNCTION_NAME, State).
+
 connected(internal, {connected, Socket, Protocol0},
 		State0=#state{owner=Owner, opts=Opts, transport=Transport}) ->
 	%% Protocol options may have been given along the protocol name.
@@ -1135,11 +1160,6 @@ connected(cast, {ws_upgrade, ReplyTo, StreamRef, _, _, _}, _) ->
 	ReplyTo ! {gun_error, self(), StreamRef, {badstate,
 		"Websocket is only supported over HTTP/1.1."}},
 	keep_state_and_data;
-connected(cast, {ws_send, Owner, Frames}, State=#state{
-		owner=Owner, protocol=Protocol=gun_ws, protocol_state=ProtoState,
-		event_handler=EvHandler, event_handler_state=EvHandlerState0}) ->
-	{Commands, EvHandlerState} = Protocol:send(Frames, ProtoState, EvHandler, EvHandlerState0),
-	commands(Commands, State#state{event_handler_state=EvHandlerState});
 connected(cast, {ws_send, ReplyTo, _}, _) ->
 	ReplyTo ! {gun_error, self(), {badstate,
 		"Connection needs to be upgraded to Websocket "
diff --git a/src/gun_ws.erl b/src/gun_ws.erl
index a6dcb00..c4eefaf 100644
--- a/src/gun_ws.erl
+++ b/src/gun_ws.erl
@@ -85,7 +85,7 @@ has_keepalive() -> false.
 init(Owner, Socket, Transport, #{stream_ref := StreamRef, headers := Headers,
 		extensions := Extensions, flow := InitialFlow, handler := Handler, opts := Opts}) ->
 	{ok, HandlerState} = Handler:init(Owner, StreamRef, Headers, Opts),
-	{connected, #ws_state{owner=Owner, stream_ref=StreamRef,
+	{connected_ws_only, #ws_state{owner=Owner, stream_ref=StreamRef,
 		socket=Socket, transport=Transport, opts=Opts, extensions=Extensions,
 		flow=InitialFlow, handler=Handler, handler_state=HandlerState}}.
 
diff --git a/test/ws_SUITE.erl b/test/ws_SUITE.erl
index c672b22..4d2387b 100644
--- a/test/ws_SUITE.erl
+++ b/test/ws_SUITE.erl
@@ -70,6 +70,16 @@ error_http10_upgrade(Config) ->
 		error(timeout)
 	end.
 
+error_http_request(Config) ->
+	doc("Ensure that requests are rejected while using Websocket."),
+	{ok, ConnPid} = gun:open("localhost", config(port, Config)),
+	{ok, _} = gun:await_up(ConnPid),
+	StreamRef1 = gun:ws_upgrade(ConnPid, "/", []),
+	{upgrade, [<<"websocket">>], _} = gun:await(ConnPid, StreamRef1),
+	StreamRef2 = gun:get(ConnPid, "/"),
+	{error, {connection_error, {badstate, _}}} = gun:await(ConnPid, StreamRef2),
+	gun:close(ConnPid).
+
 reject_upgrade(Config) ->
 	doc("Ensure Websocket connections can be rejected."),
 	{ok, ConnPid} = gun:open("localhost", config(port, Config)),