escalus

%%%===================================================================
%%% @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,
         upgrade_to_tls/2,
         use_zlib/2,
         get_transport/1,
         reset_parser/1,
         stop/1,
         kill/1,
         set_filter_predicate/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, filter_pred}).

%%%===================================================================
%%% API
%%%===================================================================

-spec connect([proplists:property()]) -> {ok, escalus:client()}.
connect(Args) ->
    {ok, Pid} = gen_server:start_link(?MODULE, [Args, self()], []),
    Transport = gen_server:call(Pid, get_transport),
    {ok, Transport}.

send(#client{rcv_pid = Pid, compress = {zlib, {_, Zout}}}, Elem) ->
    gen_server:cast(Pid, {send_compressed, Zout, Elem});
send(#client{rcv_pid = Pid}, Elem) ->
    gen_server:cast(Pid, {send, exml:to_iolist(Elem)}).

is_connected(#client{rcv_pid = Pid}) ->
    erlang:is_process_alive(Pid).

reset_parser(#client{rcv_pid = Pid}) ->
    gen_server:cast(Pid, reset_parser).

stop(#client{rcv_pid = Pid}) ->
    try
        gen_server:call(Pid, stop)
    catch
        exit:{noproc, {gen_server, call, _}} ->
            already_stopped
    end.

kill(Transport) ->
    error({not_implemented_for, ?MODULE}, [Transport]).

-spec set_filter_predicate(escalus_connection:client(),
    escalus_connection:filter_pred()) -> ok.
set_filter_predicate(#client{rcv_pid = Pid}, Pred) ->
    gen_server:call(Pid, {set_filter_pred, Pred}).


upgrade_to_tls(_, _) ->
    throw(starttls_not_supported).

%% TODO: this is en exact duplicate of escalus_tcp:use_zlib/2, DRY!
use_zlib(#client{rcv_pid = Pid} = Client, Props) ->
    escalus_connection:send(Client, escalus_stanza:compress(<<"zlib">>)),
    Compressed = escalus_connection:get_stanza(Client, compressed),
    escalus:assert(is_compressed, Compressed),
    gen_server:call(Pid, use_zlib),
    Client1 = get_transport(Client),
    {Props2, _} = escalus_session:start_stream(Client1, Props),
    {Client1, Props2}.

get_transport(#client{rcv_pid = Pid}) ->
    gen_server:call(Pid, get_transport).

%%%===================================================================
%%% gen_server callbacks
%%%===================================================================

init([Args, Owner]) ->
    Host = get_host(Args, "localhost"),
    Port = get_port(Args, 5280),
    Resource = get_resource(Args, "/ws-xmpp"),
    LegacyWS = get_legacy_ws(Args, false),
    EventClient = proplists:get_value(event_client, Args),
    SSL = proplists:get_value(ssl, Args, false),
    WSOptions = [{ssl, SSL}],
    {ok, Socket} = wsecli:start(Host, Port, Resource, WSOptions),
    Pid = self(),
    wsecli:on_open(Socket, fun() -> Pid ! opened end),
    wsecli:on_error(Socket, fun(Reason) -> Pid ! {error, Reason} end),
    wsecli:on_message(Socket, fun(Type, Data) -> Pid ! {Type, Data} end),
    wsecli:on_close(Socket, fun(_) -> Pid ! tcp_closed end),
    wait_for_socket_start(),
    ParserOpts = if
                     LegacyWS -> [];
                     true -> [{infinite_stream, true}, {autoreset, true}]
                 end,
    {ok, Parser} = exml_stream:new_parser(ParserOpts),
    {ok, #state{owner = Owner,
                socket = Socket,
                parser = Parser,
                legacy_ws = LegacyWS,
                event_client = EventClient}}.

handle_call(get_transport, _From, State) ->
    {reply, transport(State), State};
handle_call(use_zlib, _, #state{parser = Parser, socket = Socket} = State) ->
    Zin = zlib:open(),
    Zout = zlib:open(),
    ok = zlib:inflateInit(Zin),
    ok = zlib:deflateInit(Zout),
    {ok, NewParser} = exml_stream:reset_parser(Parser),
    {reply, Socket, 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(stop, _From, #state{socket = Socket,
                                compress = Compress} = State) ->
    StreamEnd = if
                    State#state.legacy_ws -> escalus_stanza:stream_end();
                    true -> escalus_stanza:ws_close()
                end,
    case Compress of
        {zlib, {Zin, Zout}} ->
            try
                ok = zlib:inflateEnd(Zin)
            catch
                error:data_error -> ok
            end,
            ok = zlib:close(Zin),
            wsecli:send(Socket, zlib:deflate(Zout,
                                             exml:to_iolist(StreamEnd),
                                             finish)),
            ok = zlib:deflateEnd(Zout),
            ok = zlib:close(Zout);
        false ->
            wsecli:send(Socket, exml:to_iolist(StreamEnd))
    end,
    wait_until_closed(),
    {stop, normal, ok, State}.

handle_cast({send_compressed, Zout, Elem}, State) ->
    wsecli:send(State#state.socket, zlib:deflate(Zout, exml:to_iolist(Elem), sync)),
    {noreply, State};
handle_cast({send, Data}, State) ->
    wsecli:send(State#state.socket, Data),
    {noreply, State};
handle_cast(reset_parser, #state{parser = Parser} = State) ->
    {ok, NewParser} = exml_stream:reset_parser(Parser),
    {noreply, State#state{parser = NewParser}}.

handle_info(tcp_closed, State) ->
    {stop, normal, State};
handle_info({error, Reason}, State) ->
    {stop, Reason, State};
handle_info({text, Data}, State) ->
    handle_data(list_to_binary(lists:flatten(Data)), State);
handle_info({binary, Data}, State) ->
    handle_data(Data, State);
handle_info(_, State) ->
    {noreply, State}.

terminate(_Reason, #state{socket = Socket} = State) ->
    common_terminate(_Reason, State),
    wsecli:stop(Socket).

code_change(_OldVsn, State, _Extra) ->
    {ok, State}.

%%%===================================================================
%%% Helpers
%%%===================================================================

handle_data(Data, State = #state{parser = Parser,
                                 compress = Compress}) ->
    {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,
    NewState = State#state{parser = NewParser},
    escalus_connection:maybe_forward_to_owner(NewState#state.filter_pred,
                                              NewState,
                                              Stanzas,
                                              fun forward_to_owner/2),
    case lists:filter(fun is_stream_end/1, Stanzas) of
        [] -> {noreply, NewState};
        _ -> {stop, normal, NewState}
    end.

-spec is_stream_end(exml_stream:element()) -> boolean().
is_stream_end(#xmlstreamend{}) -> true;
is_stream_end(_) -> false.

forward_to_owner(Stanzas, #state{owner = Owner,
                                 event_client = EventClient} = NewState) ->
    lists:foreach(fun(Stanza) ->
                          escalus_event:incoming_stanza(EventClient, Stanza),
                          Owner ! {stanza, transport(NewState), Stanza}
                  end, Stanzas).

common_terminate(_Reason, #state{parser = Parser}) ->
    exml_stream:free_parser(Parser).

transport(#state{socket = Socket,
                 compress = Compress,
                 event_client = EventClient}) ->
    #client{module = ?MODULE,
               socket = Socket,
               ssl = undefined,
               compress = Compress,
               rcv_pid = self(),
               event_client = EventClient}.

wait_until_closed() ->
    receive
        tcp_closed ->
            ok
    after ?WAIT_FOR_SOCKET_CLOSE_TIMEOUT ->
            ok
    end.

wait_for_socket_start() ->
    receive
        opened ->
            ok
    after ?HANDSHAKE_TIMEOUT ->
            throw(handshake_timeout)
    end.

-spec get_port(list(), inet:port_number()) -> inet:port_number().
get_port(Args, Default) ->
    get_option(port, Args, Default).

-spec get_host(list(), string()) -> string().
get_host(Args, Default) ->
    maybe_binary_to_list(get_option(host, Args, Default)).

-spec get_resource(list(), string()) -> string().
get_resource(Args, Default) ->
    maybe_binary_to_list(get_option(wspath, Args, Default)).

-spec get_legacy_ws(list(), boolean()) -> boolean().
get_legacy_ws(Args, Default) ->
    get_option(wslegacy, Args, Default).

-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.

-spec get_option(any(), list(), any()) -> any().
get_option(Key, Opts, Default) ->
    case lists:keyfind(Key, 1, Opts) of
        false -> Default;
        {Key, Value} -> Value
    end.