Current section
Files
Jump to
Current section
Files
src/escalus_ws.erl
%%%===================================================================
%%% @copyright (C) 2012, Erlang Solutions Ltd.
%%% @doc Module abstracting Websockets over TCP connection to XMPP server
%%% @end
%%%===================================================================
-module(escalus_ws).
-behaviour(gen_server).
-behaviour(escalus_connection).
-include_lib("exml/include/exml_stream.hrl").
-include("escalus.hrl").
%% API exports
-export([connect/1,
send/2,
is_connected/1,
reset_parser/1,
get_sm_h/1,
set_sm_h/2,
use_zlib/1,
upgrade_to_tls/2,
set_filter_predicate/2,
stop/1,
kill/1]).
%% Connection stream start and end callbacks
-export([stream_start_req/1,
stream_end_req/1,
assert_stream_start/2,
assert_stream_end/2]).
%% gen_server callbacks
-export([init/1,
handle_call/3,
handle_cast/2,
handle_info/2,
terminate/2,
code_change/3]).
-define(WAIT_FOR_SOCKET_CLOSE_TIMEOUT, 200).
-define(HANDSHAKE_TIMEOUT, 3000).
-define(SERVER, ?MODULE).
-record(state, {owner, socket, parser, legacy_ws, compress = false,
event_client, sm_state, filter_pred, stream_ref, sent_stanzas = []}).
-type state() :: #state{}.
-type sm_state() :: {boolean(), non_neg_integer(), 'active'|'inactive'}.
-type opts() :: #{
host => string(),
port => pos_integer(),
wspath => string(),
wslegacy => boolean(),
event_client => undefined | escalus_event:event_client(),
ssl => boolean(),
ssl_opts => [ssl:tls_option()],
ws_upgrade_timeout => pos_integer(),
stream_management => boolean(),
manual_ack => boolean()
}.
%%%===================================================================
%%% API
%%%===================================================================
-spec connect([proplists:property()]) -> pid().
connect(Opts0) ->
Opts1 = opts_to_map(Opts0),
{ok, Pid} = gen_server:start_link(?MODULE, [Opts1, self()], []),
Pid.
-spec send(pid(), exml:element()) -> ok.
send(Pid, Elem) ->
gen_server:cast(Pid, {send, Elem}).
-spec is_connected(pid()) -> boolean().
is_connected(Pid) ->
erlang:is_process_alive(Pid).
-spec get_sm_h(pid()) -> non_neg_integer().
get_sm_h(Pid) ->
gen_server:call(Pid, get_sm_h).
-spec set_sm_h(pid(), non_neg_integer()) -> {ok, non_neg_integer()}.
set_sm_h(Pid, H) ->
gen_server:call(Pid, {set_sm_h, H}).
-spec reset_parser(pid()) -> ok.
reset_parser(Pid) ->
gen_server:cast(Pid, reset_parser).
-spec stop(pid()) -> ok | already_stopped.
stop(Pid) ->
try
gen_server:call(Pid, stop)
catch
exit:{noproc, {gen_server, call, _}} ->
already_stopped;
exit:{normal, {gen_server, call, _}} ->
already_stopped
end.
-spec kill(pid()) -> ok | already_stopped.
kill(Pid) ->
%% Use `kill_connection` to avoid confusion with exit reason `kill`.
try
gen_server:call(Pid, kill_connection)
catch
exit:{noproc, {gen_server, call, _}} ->
already_stopped;
exit:{normal, {gen_server, call, _}} ->
already_stopped
end.
-spec set_filter_predicate(pid(), escalus_connection:filter_pred()) -> ok.
set_filter_predicate(Pid, Pred) ->
gen_server:call(Pid, {set_filter_pred, Pred}).
-spec upgrade_to_tls(_, _) -> no_return().
upgrade_to_tls(_, _) ->
throw(starttls_not_supported).
-spec use_zlib(pid()) -> ok.
use_zlib(Pid) ->
gen_server:call(Pid, use_zlib).
-spec stream_start_req(escalus_users:user_spec()) -> exml_stream:element().
stream_start_req(Props) ->
{server, Server} = lists:keyfind(server, 1, Props),
case proplists:get_value(wslegacy, Props, false) of
true ->
NS = proplists:get_value(stream_ns, Props, <<"jabber:client">>),
escalus_stanza:stream_start(Server, NS);
false ->
escalus_stanza:ws_open(Server)
end.
-spec stream_end_req(_) -> exml_stream:element().
stream_end_req(Props) ->
case proplists:get_value(wslegacy, Props, false) of
true -> escalus_stanza:stream_end();
false -> escalus_stanza:ws_close()
end.
-spec assert_stream_start(exml_stream:element(), _) -> exml_stream:element().
assert_stream_start(StreamStartRep, Props) ->
case {StreamStartRep, proplists:get_value(wslegacy, Props, false)} of
{#xmlel{name = <<"open">>}, false} ->
StreamStartRep;
{#xmlel{name = <<"open">>}, true} ->
error("<open/> with legacy WebSocket",
[StreamStartRep]);
{#xmlstreamstart{}, false} ->
error("<stream:stream> with non-legacy WebSocket",
[StreamStartRep]);
{#xmlstreamstart{}, _} ->
StreamStartRep;
_ ->
error("Not a valid stream start", [StreamStartRep])
end.
-spec assert_stream_end(exml_stream:element(), _) -> exml_stream:element().
assert_stream_end(StreamEndRep, Props) ->
case {StreamEndRep, proplists:get_value(wslegacy, Props, false)} of
{#xmlel{name = <<"close">>}, false} ->
StreamEndRep;
{#xmlel{name = <<"close">>}, true} ->
error("<close/> with legacy WebSocket",
[StreamEndRep]);
{#xmlstreamend{}, false} ->
error("<stream:stream> with non-legacy WebSocket",
[StreamEndRep]);
{#xmlstreamend{}, _} ->
StreamEndRep;
_ ->
error("Not a valid stream end", [StreamEndRep])
end.
%%%===================================================================
%%% Default options
%%%===================================================================
default_options() ->
#{host => "localhost",
port => 5280,
wspath => "/ws-xmpp",
wslegacy => false,
event_client => undefined,
ssl => false,
ssl_opts => [],
ws_upgrade_timeout => 5000,
stream_management => false,
manual_ack => false}.
%%%===================================================================
%%% gen_server callbacks
%%%===================================================================
-spec init(list()) -> {ok, state()}.
init([Opts, Owner]) ->
Opts1 = overwrite_default_opts(Opts, default_options()),
#{wspath := Resource,
wslegacy := LegacyWS,
event_client := EventClient,
ws_upgrade_timeout := Timeout} = Opts1,
SM = get_stream_management_opt(Opts1),
Resource1 = maybe_binary_to_list(Resource),
ConnPid = do_connect(Opts1),
%% Disable http2 in protocols
WSUpgradeHeaders = [{<<"sec-websocket-protocol">>, <<"xmpp">>}],
StreamRef = gun:ws_upgrade(ConnPid, Resource1, WSUpgradeHeaders,
#{protocols => [{<<"xmpp">>, gun_ws_h}]}),
wait_for_ws_upgrade(ConnPid, StreamRef, Timeout),
ParserOpts = case LegacyWS of
true -> [];
_ -> [{infinite_stream, true}, {autoreset, true}]
end,
{ok, Parser} = exml_stream:new_parser(ParserOpts),
{ok, #state{owner = Owner,
socket = ConnPid,
parser = Parser,
legacy_ws = LegacyWS,
sm_state = SM,
event_client = EventClient,
stream_ref = StreamRef}}.
wait_for_ws_upgrade(ConnPid, StreamRef, Timeout) ->
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 %% More clauses here as needed.
Timeout ->
exit(ws_upgrade_timeout)
end.
-spec handle_call(term(), {pid(), term()}, state()) ->
{reply, term(), state()} | {stop, normal, ok, state()}.
handle_call(get_sm_h, _From, #state{sm_state = {_, H, _}} = State) ->
{reply, H, State};
handle_call({set_sm_h, H}, _From, #state{sm_state = {A, _OldH, S}} = State) ->
NewState = State#state{sm_state={A, H, S}},
{reply, {ok, H}, NewState};
handle_call(use_zlib, _, #state{parser = Parser} = State) ->
Zin = zlib:open(),
Zout = zlib:open(),
ok = zlib:inflateInit(Zin),
ok = zlib:deflateInit(Zout),
{ok, NewParser} = exml_stream:reset_parser(Parser),
{reply, ok, State#state{parser = NewParser,
compress = {zlib, {Zin, Zout}}}};
handle_call({set_filter_pred, Pred}, _From, State) ->
{reply, ok, State#state{filter_pred = Pred}};
handle_call(kill_connection, _, #state{socket = ConnPid} = State) ->
gun:close(ConnPid),
{stop, normal, ok, State};
handle_call(stop, _From, #state{socket = ConnPid, stream_ref = StreamRef} = State) ->
close_compression_streams(State#state.compress),
gun:ws_send(ConnPid, StreamRef, close),
{stop, normal, ok, State}.
-spec handle_cast(term(), state()) -> {noreply, state()}.
handle_cast({send, Elem}, State) ->
Data = case State#state.compress of
{zlib, {_, Zout}} -> zlib:deflate(Zout, exml:to_iolist(Elem), sync);
false -> exml:to_iolist(Elem)
end,
gun:ws_send(State#state.socket, State#state.stream_ref, {text, Data}),
{noreply, State};
handle_cast(reset_parser, #state{parser = Parser} = State) ->
{ok, NewParser} = exml_stream:reset_parser(Parser),
{noreply, State#state{parser = NewParser}}.
-spec handle_info(term(), state()) -> {noreply, state()} | {stop, term(), state()}.
handle_info(tcp_closed, State) ->
{stop, normal, State};
handle_info({error, Reason}, State) ->
{stop, Reason, State};
handle_info({gun_ws, ConnPid, _StreamRef, {text, Data}}, #state{socket = ConnPid} = State) ->
handle_data(Data, State);
handle_info(_, State) ->
{noreply, State}.
-spec terminate(term(), state()) -> term().
terminate(Reason, State) ->
common_terminate(Reason, State).
-spec code_change(term(), state(), term()) -> {ok, state()}.
code_change(_OldVsn, State, _Extra) ->
{ok, State}.
%%%===================================================================
%%% Helpers
%%%===================================================================
-spec get_stream_management_opt(opts()) -> sm_state().
get_stream_management_opt(#{stream_management := false}) ->
{false, 0, inactive};
get_stream_management_opt(#{manual_ack := true}) ->
{false, 0, inactive};
get_stream_management_opt(#{stream_management := true, manual_ack := false}) ->
{true, 0, inactive}.
overwrite_default_opts(GivenOpts, DefaultOpts) ->
maps:merge(DefaultOpts, GivenOpts).
do_connect(#{ssl := true, ssl_opts := SSLOpts} = Opts) ->
TransportOpts = #{transport => tls, protocols => [http],
tls_opts => SSLOpts},
do_connect(Opts, TransportOpts);
do_connect(Opts) ->
do_connect(Opts, #{transport => tcp, protocols => [http]}).
do_connect(#{host := Host, port := Port}, TransportOpts) ->
Host1 = maybe_binary_to_list(Host),
{ok, ConnPid} = gun:open(Host1, Port, TransportOpts),
{ok, http} = gun:await_up(ConnPid),
ConnPid.
handle_data(Data, State = #state{parser = Parser,
compress = Compress}) ->
Timestamp = os:system_time(microsecond),
{ok, NewParser, Stanzas} =
case Compress of
false ->
exml_stream:parse(Parser, Data);
{zlib, {Zin,_}} ->
Decompressed = iolist_to_binary(zlib:inflate(Zin, Data)),
exml_stream:parse(Parser, Decompressed)
end,
FwdState = State#state{parser = NewParser, sent_stanzas = []},
NewState = escalus_connection:maybe_forward_to_owner(FwdState#state.filter_pred,
FwdState,
Stanzas,
fun forward_to_owner/3, Timestamp),
case lists:filter(fun(Stanza) -> is_stream_end(Stanza, State) end, Stanzas) of
[] -> {noreply, NewState};
_ -> {stop, normal, NewState}
end.
-spec is_stream_end(exml_stream:element(), state()) -> boolean().
is_stream_end(#xmlstreamend{}, #state{legacy_ws = true}) -> true;
is_stream_end(#xmlel{name = <<"close">>}, #state{legacy_ws = false}) -> true;
is_stream_end(_, _) -> false.
forward_to_owner(Stanzas0, #state{owner = Owner,
sm_state = SM0,
event_client = EventClient} = State, Timestamp) ->
{SM1, AckRequests, StanzasNoRs} = escalus_connection:separate_ack_requests(SM0, Stanzas0),
reply_to_ack_requests(SM1, AckRequests, State),
lists:foreach(fun(Stanza) ->
escalus_event:incoming_stanza(EventClient, Stanza),
Owner ! escalus_connection:stanza_msg(Stanza, #{recv_timestamp => Timestamp})
end, StanzasNoRs),
State#state{sm_state = SM1, sent_stanzas = StanzasNoRs}.
reply_to_ack_requests({false, H, A}, _, _) ->
{false, H, A};
reply_to_ack_requests({true, H, inactive}, _, _) ->
{true, H, inactive};
reply_to_ack_requests({true, H0, active}, Acks, State) ->
{true,
lists:foldl(fun({Ack, H}, _) ->
Ack1 = exml:to_iolist(Ack),
gun:ws_send(State#state.socket, State#state.stream_ref, {text, Ack1}),
H
end, H0, Acks),
active}.
common_terminate(_Reason, #state{parser = Parser}) ->
exml_stream:free_parser(Parser).
-spec maybe_binary_to_list(binary() | string()) -> string().
maybe_binary_to_list(B) when is_binary(B) -> binary_to_list(B);
maybe_binary_to_list(S) when is_list(S) -> S.
close_compression_streams(false) ->
ok;
close_compression_streams({zlib, {Zin, Zout}}) ->
try
zlib:deflate(Zout, <<>>, finish),
ok = zlib:inflateEnd(Zin),
ok = zlib:deflateEnd(Zout)
catch
error:data_error -> ok
after
ok = zlib:close(Zin),
ok = zlib:close(Zout)
end.
-spec opts_to_map(proplists:proplist() | opts()) -> opts().
opts_to_map(Opts) when is_map(Opts) -> Opts;
opts_to_map(Opts) when is_list(Opts) -> maps:from_list(Opts).