diff options
| -rw-r--r-- | src/gun.erl | 34 | ||||
| -rw-r--r-- | src/gun_ws.erl | 2 | ||||
| -rw-r--r-- | test/ws_SUITE.erl | 10 |
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)), |