escalus

%%%===================================================================
%%% @copyright (C) 2011-2012, Erlang Solutions Ltd.
%%% @doc Module abstracting TCP connection to XMPP server
%%% @end
%%%===================================================================

-module(escalus_tcp).
-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,
         get_sm_h/1,
         set_sm_h/2,
         set_filter_predicate/2,
         stop/1,
         kill/1]).

%% gen_server callbacks
-export([init/1,
         handle_call/3,
         handle_cast/2,
         handle_info/2,
         terminate/2,
         code_change/3]).

%% Low level API
-export([get_active/1,
         set_active/2,
         recv/1]).

-ifdef(EUNIT_TEST).
-compile(export_all).
-endif.

%% Stream management automation
%%               :: {Auto-ack?,                H,         counting Hs?}.
-type sm_state() :: {boolean(), non_neg_integer(), 'active'|'inactive'}.

-export_type([sm_state/0]).

-define(WAIT_FOR_SOCKET_CLOSE_TIMEOUT, 200).
-define(SERVER, ?MODULE).
-include("escalus_tcp.hrl").

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

-spec connect([proplists:property()]) -> {ok, #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} = Client, Elem) ->
    gen_server:cast(Pid, {send, Client, Elem}).

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

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

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

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

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

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

kill(#client{rcv_pid = Pid}) ->
    %% Use `kill_connection` to avoid confusion with exit reason `kill`.
    gen_server:call(Pid, kill_connection).

upgrade_to_tls(#client{socket = Socket, rcv_pid = Pid} = Client, Props) ->
    Starttls = escalus_stanza:starttls(),
    gen_tcp:send(Socket, exml:to_iolist(Starttls)),
    escalus_connection:get_stanza(Client, proceed),
    SSLOpts = proplists:get_value(ssl_opts, Props, []),
    case gen_server:call(Pid, {upgrade_to_tls, SSLOpts}) of
        {error, Error} ->
            error(Error);
        _ ->
            Client2 = get_transport(Client),
            {Props2, _} = escalus_session:start_stream(Client2, Props),
            {Client2, Props2}
    end.

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

%%%===================================================================
%%% Low level API
%%%===================================================================

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

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

-spec recv(#client{}) -> exml_stream:element() | empty.
recv(#client{rcv_pid = Pid}) ->
    gen_server:call(Pid, recv).

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

init([Args, Owner]) ->
    Host = proplists:get_value(host, Args, <<"localhost">>),
    Port = proplists:get_value(port, Args, 5222),

    Address = host_to_inet(Host),
    EventClient = proplists:get_value(event_client, Args),
    Interface = proplists:get_value(iface, Args),
    IsSSLConnection = proplists:get_value(ssl, Args, false),

    OnReplyFun = proplists:get_value(on_reply, Args, fun(_) -> ok end),
    OnRequestFun = proplists:get_value(on_request, Args, fun(_) -> ok end),
    OnConnectFun = proplists:get_value(on_connect, Args, fun(_) -> ok end),

    SM = case {proplists:get_value(stream_management, Args, false),
               proplists:get_value(manual_ack, Args, false)}
         of
             {false,_}          -> {false, 0, inactive};
             {_, true}          -> {false, 0, inactive};
             {true,false}       -> {true, 0, inactive}
         end,


    BasicOpts = [binary, {active, once}],
    SocketOpts = case Interface of
                     undefined -> BasicOpts;
                     _         -> [{ip, iface_to_ip_address(Interface)}] ++ BasicOpts
                 end,

    {ok, Socket} = do_connect(IsSSLConnection, Address, Port, Args,
                              SocketOpts, OnConnectFun),
    {ok, Parser} = exml_stream:new_parser(),
    {ok, #state{owner = Owner,
                socket = Socket,
                parser = Parser,
                ssl = IsSSLConnection,
                sm_state = SM,
                event_client = EventClient,
                on_reply = OnReplyFun,
                on_request = OnRequestFun}}.

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(get_transport, _From, State) ->
    {reply, transport(State), State};
handle_call({upgrade_to_tls, SSLOpts}, _From, #state{socket = Socket} = State) ->
    SSLOpts1 = [{reuse_sessions, true}],
    SSLOpts2 = lists:keymerge(1, lists:keysort(1, SSLOpts),
                              lists:keysort(1, SSLOpts1)),
    case ssl:connect(Socket, SSLOpts2) of
        {ok, Socket2} ->
            {ok, Parser} = exml_stream:new_parser(),
            {reply, Socket2,
             State#state{socket = Socket2, parser = Parser, ssl=true}};
        {error, closed} = E ->
            {reply, E, State}
    end;
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(get_active, _From, #state{active = Active} = State) ->
    {reply, Active, State};
handle_call({set_active, Active}, _From, State) ->
    {reply, ok, State#state{active = Active}};
handle_call({set_filter_pred, Pred}, _From, State) ->
    {reply, ok, State#state{filter_pred = Pred}};
handle_call(recv, _From, State) ->
    {Reply, NS} = handle_recv(State),
    {reply, Reply, NS};
handle_call(kill_connection, _, #state{socket = Socket } = S) ->
    gen_tcp:close(Socket),
    close_compression_streams(S#state.compress),
    {stop, normal, ok, S};
handle_call(stop, _From, #state{} = S) ->
    send_stream_end(S),
    close_compression_streams(S#state.compress),
    wait_until_closed(S#state.socket),
    {stop, normal, ok, S}.

handle_cast({send, #client{socket = Socket, ssl = Ssl, compress = Compress},
             Elem}, #state{on_request = OnRequestFun} = State) ->
    Reply = case {Ssl, Compress} of
                {true, _} ->
                    ssl:send(Socket, exml:to_iolist(Elem));
                {false, {zlib, {_,Zout}}} ->
                    Deflated = zlib:deflate(Zout, exml:to_iolist(Elem), sync),
                    gen_tcp:send(State#state.socket, Deflated);
                {false, false} ->
                    gen_tcp:send(Socket, exml:to_iolist(Elem))
            end,
    OnRequestFun(Reply),
    {noreply, State};
handle_cast(reset_parser, #state{parser = Parser} = State) ->
    {ok, NewParser} = exml_stream:reset_parser(Parser),
    {noreply, State#state{parser = NewParser}};
handle_cast(stop, State) ->
    {stop, normal, State}.

handle_info({tcp, Socket, Data}, State) ->
    inet:setopts(Socket, [{active, once}]),
    NewState = handle_data(Socket, Data, State),
    {noreply, NewState};
handle_info({ssl, Socket, Data}, State) ->
    ssl:setopts(Socket, [{active, once}]),
    NewState = handle_data(Socket, Data, State),
    {noreply, NewState};
handle_info({tcp_closed, Socket}, #state{socket = Socket} = State) ->
    {stop, normal, State};
handle_info(_, State) ->
    {noreply, State}.

terminate(_Reason, #state{socket = Socket, ssl = true} = State) ->
    common_terminate(_Reason, State),
    ssl:close(Socket);
terminate(_Reason, #state{socket = Socket} = State) ->
    common_terminate(_Reason, State),
    gen_tcp:close(Socket).

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

%%%===================================================================
%%% Helpers
%%%===================================================================
handle_data(Socket, Data, #state{parser = Parser,
                                 socket = Socket,
                                 compress = Compress,
                                 on_reply = OnReplyFun} = State) ->
    OnReplyFun({erlang:byte_size(Data)}),
    {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},
    case State#state.active of
        true ->
            escalus_connection:maybe_forward_to_owner(NewState#state.filter_pred,
                                                      NewState, Stanzas,
                                                      fun forward_to_owner/2);
        false ->
            store_reply(Stanzas, NewState)
    end.


forward_to_owner(Stanzas0, #state{owner = Owner,
                                  sm_state = SM0,
                                  event_client = EventClient} = State) ->
    {SM1, AckRequests, StanzasNoRs} = separate_ack_requests(SM0, Stanzas0),
    reply_to_ack_requests(SM1, AckRequests, State),
    NewState = State#state{sm_state=SM1},

    lists:foreach(fun(Stanza) ->
        escalus_event:incoming_stanza(EventClient, Stanza),
        Owner ! {stanza, transport(NewState), Stanza}
    end, StanzasNoRs),

    case lists:keyfind(xmlstreamend, 1, StanzasNoRs) of
        false -> ok;
        _     -> gen_server:cast(self(), stop)
    end,
    NewState#state{replies = StanzasNoRs}.

store_reply(Stanzas, #state{replies = Replies} = S) ->
    S#state{replies = Replies ++ Stanzas}.

%% @doc this looks like it's only used in esl/escalus_tests bosh_SUITE
%% Maybe it can be removed?
handle_recv(#state{replies = []} = S) ->
    {empty, S};
handle_recv(#state{replies = [Reply | Replies]} = S) ->
    case Reply of
        #xmlstreamend{} ->
            gen_server:cast(self(), stop);
        _ -> ok
    end,
    {Reply, S#state{replies = Replies}}.

separate_ack_requests({false, H0, A}, Stanzas) ->
    %% Don't keep track of H
    {{false, H0, A}, [], Stanzas};
separate_ack_requests({true, H0, inactive}, Stanzas) ->
    Enabled = [ S || S <- Stanzas, escalus_pred:is_sm_enabled(S)],
    Resumed = [ S || S <- Stanzas, escalus_pred:is_sm_resumed(S)],

    case {length(Enabled),length(Resumed)} of
        %% Enabled SM: set the H param to 0 and activate counter.
        {1,0} -> {{true, 0, active}, [], Stanzas};

        %% Resumed SM: keep the H param and activate counter.
        {0,1} -> {{true, H0, active}, [], Stanzas};

        %% No new SM state: continue as usual
        {0,0} -> {{true, H0, inactive}, [], Stanzas}
    end;
separate_ack_requests({true, H0, active}, Stanzas) ->
    %% Count H and construct appropriate acks
    F = fun(Stanza, {H, Acks, NonAckRequests}) ->
                case escalus_pred:is_sm_ack_request(Stanza) of
                    true -> {H, [make_ack(H)|Acks], NonAckRequests};
                    false -> {H+1, Acks, [Stanza|NonAckRequests]}
                end
        end,
    {H, Acks, Others} = lists:foldl(F, {H0, [], []}, Stanzas),
    {{true, H, active}, lists:reverse(Acks), lists:reverse(Others)}.

make_ack(H) -> {escalus_stanza:sm_ack(H), H}.

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}, _) -> raw_send(State, Ack), H end,
                 H0, Acks),
     active}.

raw_send(#state{socket=Socket, ssl=true}, Elem) ->
    ssl:send(Socket, exml:to_iolist(Elem));
raw_send(#state{socket=_Socket, compress=true}, _Elem) ->
    throw({escalus_tcp, auto_ack_not_implemented_for_compressed_streams});
raw_send(#state{socket=Socket}, Elem) ->
    gen_tcp:send(Socket, exml:to_iolist(Elem)).

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

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

wait_until_closed(Socket) ->
    receive
        {tcp_closed, Socket} ->
            ok
    after ?WAIT_FOR_SOCKET_CLOSE_TIMEOUT ->
            ok
    end.

-spec host_to_inet(tuple() | atom() | list() | binary())
    -> inet:ip_address() | inet:hostname().
host_to_inet({_,_,_,_} = IP4) -> IP4;
host_to_inet({_,_,_,_,_,_,_,_} = IP6) -> IP6;
host_to_inet(Address) when is_list(Address) orelse is_atom(Address) -> Address;
host_to_inet(BAddress) when is_binary(BAddress) -> binary_to_list(BAddress).

iface_to_ip_address({_,_,_,_} = IP4) -> IP4;
iface_to_ip_address({_,_,_,_,_,_,_,_} = IP6) -> IP6.

close_compression_streams(false) ->
    ok;
close_compression_streams({zlib, {Zin, Zout}}) ->
    try
        ok = zlib:inflateEnd(Zin),
        ok = zlib:deflateEnd(Zout)
    catch
        error:data_error -> ok
    after
        ok = zlib:close(Zin),
        ok = zlib:close(Zout)
    end.

send_stream_end(#state{socket = Socket, ssl = Ssl, compress = Compress}) ->
    StreamEnd = escalus_stanza:stream_end(),
    case {Ssl, Compress} of
        {true, _} ->
            ssl:send(Socket, exml:to_iolist(StreamEnd));
        {false, {zlib, {_, Zout}}} ->
            gen_tcp:send(Socket, zlib:deflate(Zout,
                                              exml:to_iolist(StreamEnd),
                                              finish));
        {false, false} ->
            gen_tcp:send(Socket, exml:to_iolist(StreamEnd))
    end.

do_connect(IsSSLConnection, Address, Port, Args, SocketOpts, OnConnectFun) ->
    TimeB = os:timestamp(),
    Reply = maybe_ssl_connection(IsSSLConnection, Address, Port, SocketOpts, Args),
    TimeA = os:timestamp(),
    ConnectionTime = timer:now_diff(TimeA, TimeB),
    case Reply of
        {ok, Socket} ->
            OnConnectFun({ok, Socket, ConnectionTime});
        {error, _} ->
            OnConnectFun(Reply)
    end,
    Reply.

maybe_ssl_connection(true, Address, Port, Opts, Args) ->
    SSLOpts = proplists:get_value(ssl_opts, Args, []),
    ssl:connect(Address, Port, Opts ++ SSLOpts);
maybe_ssl_connection(_, Address, Port, Opts, _) ->
    gen_tcp:connect(Address, Port, Opts).