Rewrite S2S and ejabberd_service code to use XML generator

1 files changed, 166 insertions, 213 deletions

diff --git a/src/ejabberd_service.erl b/src/ejabberd_service.erl
index 465fb587..432253e0 100644
--- a/src/ejabberd_service.erl
+++ b/src/ejabberd_service.erl
@@ -46,8 +46,7 @@
 
 -include("ejabberd.hrl").
 -include("logger.hrl").
-
--include("jlib.hrl").
+-include("xmpp.hrl").
 
 -record(state,
 	{socket                    :: ejabberd_socket:socket_state(),
@@ -58,44 +57,19 @@
          access                    :: atom(),
 	 check_from = true         :: boolean()}).
 
-%-define(DBGFSM, true).
+-type state_name() :: wait_for_stream | wait_for_handshake | stream_established.
+-type state() :: #state{}.
+-type fsm_next() :: {next_state, state_name(), state()}.
+-type fsm_stop() :: {stop, normal, state()}.
+-type fsm_transition() :: fsm_stop() | fsm_next().
 
+%-define(DBGFSM, true).
 -ifdef(DBGFSM).
-
 -define(FSMOPTS, [{debug, [trace]}]).
-
 -else.
-
 -define(FSMOPTS, []).
-
 -endif.
 
--define(STREAM_HEADER,
-	<<"<?xml version='1.0'?><stream:stream "
-	  "xmlns:stream='http://etherx.jabber.org/stream"
-	  "s' xmlns='jabber:component:accept' id='~s' "
-	  "from='~s'>">>).
-
--define(STREAM_TRAILER, <<"</stream:stream>">>).
-
--define(INVALID_HEADER_ERR,
-	<<"<stream:stream xmlns:stream='http://etherx.ja"
-	  "bber.org/streams'><stream:error>Invalid "
-	  "Stream Header</stream:error></stream:stream>">>).
-
--define(INVALID_HANDSHAKE_ERR,
-	<<"<stream:error><not-authorized xmlns='urn:ietf"
-	  ":params:xml:ns:xmpp-streams'/><text "
-	  "xmlns='urn:ietf:params:xml:ns:xmpp-streams' "
-	  "xml:lang='en'>Invalid Handshake</text></strea"
-	  "m:error></stream:stream>">>).
-
--define(INVALID_XML_ERR,
-	fxml:element_to_binary(?SERR_XML_NOT_WELL_FORMED)).
-
--define(INVALID_NS_ERR,
-	fxml:element_to_binary(?SERR_INVALID_NAMESPACE)).
-
 %%%----------------------------------------------------------------------
 %%% API
 %%%----------------------------------------------------------------------
@@ -112,14 +86,6 @@ socket_type() -> xml_stream.
 %%%----------------------------------------------------------------------
 %%% Callback functions from gen_fsm
 %%%----------------------------------------------------------------------
-
-%%----------------------------------------------------------------------
-%% Func: init/1
-%% Returns: {ok, StateName, StateData}          |
-%%          {ok, StateName, StateData, Timeout} |
-%%          ignore                              |
-%%          {stop, StopReason}
-%%----------------------------------------------------------------------
 init([{SockMod, Socket}, Opts]) ->
     ?INFO_MSG("(~w) External service connected", [Socket]),
     Access = case lists:keysearch(access, 1, Opts) of
@@ -157,177 +123,127 @@ init([{SockMod, Socket}, Opts]) ->
 	    streamid = new_id(), host_opts = HostOpts,
 	    access = Access, check_from = CheckFrom}}.
 
-%%----------------------------------------------------------------------
-%% Func: StateName/2
-%% Returns: {next_state, NextStateName, NextStateData}          |
-%%          {next_state, NextStateName, NextStateData, Timeout} |
-%%          {stop, Reason, NewStateData}
-%%----------------------------------------------------------------------
-
-wait_for_stream({xmlstreamstart, _Name, Attrs},
-		StateData) ->
-    case fxml:get_attr_s(<<"xmlns">>, Attrs) of
-      <<"jabber:component:accept">> ->
-	  To = fxml:get_attr_s(<<"to">>, Attrs),
-	  Host = jid:nameprep(To),
-	  if Host == error ->
-		  Header = io_lib:format(?STREAM_HEADER,
-					 [<<"none">>, ?MYNAME]),
-		  send_text(StateData,
-			    <<(list_to_binary(Header))/binary,
-			      (?INVALID_XML_ERR)/binary,
-			      (?STREAM_TRAILER)/binary>>),
-		  {stop, normal, StateData};
-	     true ->
-		  Header = io_lib:format(?STREAM_HEADER,
-					 [StateData#state.streamid, fxml:crypt(To)]),
-		  send_text(StateData, Header),
-		  HostOpts = case dict:is_key(Host, StateData#state.host_opts) of
-				 true ->
-				     StateData#state.host_opts;
-				 false ->
-				     case dict:find(global, StateData#state.host_opts) of
-					 {ok, GlobalPass} ->
-					     dict:from_list([{Host, GlobalPass}]);
-					 error ->
-					     StateData#state.host_opts
-				     end
-			     end,
-		  {next_state, wait_for_handshake,
-		   StateData#state{host = Host, host_opts = HostOpts}}
-	  end;
-      _ ->
-	  send_text(StateData, ?INVALID_HEADER_ERR),
-	  {stop, normal, StateData}
+wait_for_stream({xmlstreamstart, Name, Attrs}, StateData) ->
+    try xmpp:decode(#xmlel{name = Name, attrs = Attrs}) of
+	#stream_start{xmlns = ?NS_COMPONENT, to = To} when is_record(To, jid) ->
+	    Host = To#jid.lserver,
+	    send_header(StateData, To),
+	    HostOpts = case dict:is_key(Host, StateData#state.host_opts) of
+			   true ->
+			       StateData#state.host_opts;
+			   false ->
+			       case dict:find(global, StateData#state.host_opts) of
+				   {ok, GlobalPass} ->
+				       dict:from_list([{Host, GlobalPass}]);
+				   error ->
+				       StateData#state.host_opts
+			       end
+		       end,
+	    {next_state, wait_for_handshake,
+	     StateData#state{host = Host, host_opts = HostOpts}};
+	#stream_start{xmlns = ?NS_COMPONENT} ->
+	    send_header(StateData, ?MYNAME),
+	    send_element(StateData, xmpp:serr_improper_addressing()),
+	    {stop, normal, StateData};
+	#stream_start{} ->
+	    send_header(StateData, ?MYNAME),
+	    send_element(StateData, xmpp:serr_invalid_namespace()),
+	    {stop, normal, StateData}
+    catch _:{xmpp_codec, Why} ->
+	    Txt = xmpp:format_error(Why),
+	    send_header(StateData, ?MYNAME),
+	    send_element(StateData, xmpp:serr_invalid_xml(Txt, ?MYLANG)),
+	    {stop, normal, StateData}
     end;
 wait_for_stream({xmlstreamerror, _}, StateData) ->
-    Header = io_lib:format(?STREAM_HEADER,
-			   [<<"none">>, ?MYNAME]),
-    send_text(StateData,
-	      <<(list_to_binary(Header))/binary, (?INVALID_XML_ERR)/binary,
-		(?STREAM_TRAILER)/binary>>),
+    send_header(StateData, ?MYNAME),
+    send_element(StateData, xmpp:serr_not_well_formed()),
     {stop, normal, StateData};
 wait_for_stream(closed, StateData) ->
     {stop, normal, StateData}.
 
 wait_for_handshake({xmlstreamelement, El}, StateData) ->
-    #xmlel{name = Name, children = Els} = El,
-    case {Name, fxml:get_cdata(Els)} of
-      {<<"handshake">>, Digest} ->
-	  case dict:find(StateData#state.host, StateData#state.host_opts) of
-	      {ok, Password} ->
-		  case p1_sha:sha(<<(StateData#state.streamid)/binary,
-				    Password/binary>>) of
-		      Digest ->
-			  send_text(StateData, <<"<handshake/>">>),
-			  lists:foreach(
-			    fun (H) ->
-				    ejabberd_router:register_route(H, ?MYNAME),
-				    ?INFO_MSG("Route registered for service ~p~n",
-					      [H])
-			    end, dict:fetch_keys(StateData#state.host_opts)),
-			  {next_state, stream_established, StateData};
-		      _ ->
-			  send_text(StateData, ?INVALID_HANDSHAKE_ERR),
-			  {stop, normal, StateData}
-		  end;
-	      _ ->
-		  send_text(StateData, ?INVALID_HANDSHAKE_ERR),
-		  {stop, normal, StateData}
-	  end;
-      _ -> {next_state, wait_for_handshake, StateData}
+    decode_element(El, wait_for_handshake, StateData);
+wait_for_handshake(#handshake{data = Digest}, StateData) ->
+    case dict:find(StateData#state.host, StateData#state.host_opts) of
+	{ok, Password} ->
+	    case p1_sha:sha(<<(StateData#state.streamid)/binary,
+			      Password/binary>>) of
+		Digest ->
+		    send_element(StateData, #handshake{}),
+		    lists:foreach(
+		      fun (H) ->
+			      ejabberd_router:register_route(H, ?MYNAME),
+			      ?INFO_MSG("Route registered for service ~p~n",
+					[H])
+		      end, dict:fetch_keys(StateData#state.host_opts)),
+		    {next_state, stream_established, StateData};
+		_ ->
+		    send_element(StateData, xmpp:serr_not_authorized()),
+		    {stop, normal, StateData}
+	    end;
+	_ ->
+	    send_element(StateData, xmpp:serr_not_authorized()),
+	    {stop, normal, StateData}
     end;
 wait_for_handshake({xmlstreamend, _Name}, StateData) ->
     {stop, normal, StateData};
 wait_for_handshake({xmlstreamerror, _}, StateData) ->
-    send_text(StateData,
-	      <<(?INVALID_XML_ERR)/binary,
-		(?STREAM_TRAILER)/binary>>),
+    send_element(StateData, xmpp:serr_not_well_formed()),
     {stop, normal, StateData};
 wait_for_handshake(closed, StateData) ->
-    {stop, normal, StateData}.
+    {stop, normal, StateData};
+wait_for_handshake(_Pkt, StateData) ->
+    {next_state, wait_for_handshake, StateData}.
 
 stream_established({xmlstreamelement, El}, StateData) ->
-    NewEl = jlib:remove_attr(<<"xmlns">>, El),
-    #xmlel{name = Name, attrs = Attrs} = NewEl,
-    From = fxml:get_attr_s(<<"from">>, Attrs),
-    FromJID = case StateData#state.check_from of
-		%% If the admin does not want to check the from field
-		%% when accept packets from any address.
-		%% In this case, the component can send packet of
-		%% behalf of the server users.
-		false -> jid:from_string(From);
-		%% The default is the standard behaviour in XEP-0114
-		_ ->
-		    FromJID1 = jid:from_string(From),
-		    case FromJID1 of
-		      #jid{lserver = Server} ->
-			  case dict:is_key(Server, StateData#state.host_opts) of
-			    true -> FromJID1;
-			    false -> error
-			  end;
-		      _ -> error
-		    end
-	      end,
-    To = fxml:get_attr_s(<<"to">>, Attrs),
-    ToJID = case To of
-	      <<"">> -> error;
-	      _ -> jid:from_string(To)
-	    end,
-    if ((Name == <<"iq">>) or (Name == <<"message">>) or
-	  (Name == <<"presence">>))
-	 and (ToJID /= error)
-	 and (FromJID /= error) ->
-	   ejabberd_router:route(FromJID, ToJID, NewEl);
+    decode_element(El, stream_established, StateData);
+stream_established(El, StateData) when ?is_stanza(El) ->
+    From = xmpp:get_from(El),
+    To = xmpp:get_to(El),
+    Lang = xmpp:get_lang(El),
+    if From == undefined orelse To == undefined ->
+	    send_error(StateData, El, xmpp:err_jid_malformed());
        true ->
-	   Lang = fxml:get_tag_attr_s(<<"xml:lang">>, El),
-	   Txt = <<"Incorrect stanza name or from/to JID">>,
-	   Err = jlib:make_error_reply(NewEl, ?ERRT_BAD_REQUEST(Lang, Txt)),
-	   send_element(StateData, Err),
-	   error
+	    FromJID = case StateData#state.check_from of
+			  false ->
+			      %% If the admin does not want to check the from field
+			      %% when accept packets from any address.
+			      %% In this case, the component can send packet of
+			      %% behalf of the server users.
+			      From;
+			  _ ->
+			      %% The default is the standard behaviour in XEP-0114
+			      case From of
+				  #jid{lserver = Server} ->
+				      case dict:is_key(Server, StateData#state.host_opts) of
+					  true -> From;
+					  false -> error
+				      end;
+				  _ -> error
+			      end
+		      end,
+	    if FromJID /= error ->
+		    ejabberd_router:route(FromJID, To, El);
+	       true ->
+		    Txt = <<"Incorrect value of 'from' or 'to' attribute">>,
+		    send_error(StateData, El, xmpp:err_not_allowed(Txt, Lang))
+	    end
     end,
     {next_state, stream_established, StateData};
 stream_established({xmlstreamend, _Name}, StateData) ->
     {stop, normal, StateData};
 stream_established({xmlstreamerror, _}, StateData) ->
-    send_text(StateData,
-	      <<(?INVALID_XML_ERR)/binary,
-		(?STREAM_TRAILER)/binary>>),
+    send_element(StateData, xmpp:serr_not_well_formed()),
     {stop, normal, StateData};
 stream_established(closed, StateData) ->
-    {stop, normal, StateData}.
-
-%%----------------------------------------------------------------------
-%% Func: StateName/3
-%% Returns: {next_state, NextStateName, NextStateData}            |
-%%          {next_state, NextStateName, NextStateData, Timeout}   |
-%%          {reply, Reply, NextStateName, NextStateData}          |
-%%          {reply, Reply, NextStateName, NextStateData, Timeout} |
-%%          {stop, Reason, NewStateData}                          |
-%%          {stop, Reason, Reply, NewStateData}
-%%----------------------------------------------------------------------
-%state_name(Event, From, StateData) ->
-%    Reply = ok,
-%    {reply, Reply, state_name, StateData}.
+    {stop, normal, StateData};
+stream_established(_Event, StateData) ->
+    {next_state, stream_established, StateData}.
 
-%%----------------------------------------------------------------------
-%% Func: handle_event/3
-%% Returns: {next_state, NextStateName, NextStateData}          |
-%%          {next_state, NextStateName, NextStateData, Timeout} |
-%%          {stop, Reason, NewStateData}
-%%----------------------------------------------------------------------
 handle_event(_Event, StateName, StateData) ->
     {next_state, StateName, StateData}.
 
-%%----------------------------------------------------------------------
-%% Func: handle_sync_event/4
-%% Returns: {next_state, NextStateName, NextStateData}            |
-%%          {next_state, NextStateName, NextStateData, Timeout}   |
-%%          {reply, Reply, NextStateName, NextStateData}          |
-%%          {reply, Reply, NextStateName, NextStateData, Timeout} |
-%%          {stop, Reason, NewStateData}                          |
-%%          {stop, Reason, Reply, NewStateData}
-%%----------------------------------------------------------------------
 handle_sync_event(_Event, _From, StateName,
 		  StateData) ->
     Reply = ok, {reply, Reply, StateName, StateData}.
@@ -335,12 +251,6 @@ handle_sync_event(_Event, _From, StateName,
 code_change(_OldVsn, StateName, StateData, _Extra) ->
     {ok, StateName, StateData}.
 
-%%----------------------------------------------------------------------
-%% Func: handle_info/3
-%% Returns: {next_state, NextStateName, NextStateData}          |
-%%          {next_state, NextStateName, NextStateData, Timeout} |
-%%          {stop, Reason, NewStateData}
-%%----------------------------------------------------------------------
 handle_info({send_text, Text}, StateName, StateData) ->
     send_text(StateData, Text),
     {next_state, StateName, StateData};
@@ -349,34 +259,20 @@ handle_info({send_element, El}, StateName, StateData) ->
     {next_state, StateName, StateData};
 handle_info({route, From, To, Packet}, StateName,
 	    StateData) ->
-    case acl:match_rule(global, StateData#state.access,
-			From)
-	of
+    case acl:match_rule(global, StateData#state.access, From) of
       allow ->
-	  #xmlel{name = Name, attrs = Attrs, children = Els} =
-	      Packet,
-	  Attrs2 =
-	      jlib:replace_from_to_attrs(jid:to_string(From),
-					 jid:to_string(To), Attrs),
-	  Text = fxml:element_to_binary(#xmlel{name = Name,
-					      attrs = Attrs2, children = Els}),
-	  send_text(StateData, Text);
-      deny ->
-	  Lang = fxml:get_tag_attr_s(<<"xml:lang">>, Packet),
-	  Txt = <<"Denied by ACL">>,
-	  Err = jlib:make_error_reply(Packet, ?ERRT_NOT_ALLOWED(Lang, Txt)),
-	  ejabberd_router:route_error(To, From, Err, Packet)
+	    Pkt = xmpp:set_from_to(Packet, From, To),
+	    send_element(StateData, Pkt);
+	deny ->
+	    Lang = xmpp:get_lang(Packet),
+	    Err = xmpp:err_not_allowed(<<"Denied by ACL">>, Lang),
+	    ejabberd_router:route_error(To, From, Packet, Err)
     end,
     {next_state, StateName, StateData};
 handle_info(Info, StateName, StateData) ->
     ?ERROR_MSG("Unexpected info: ~p", [Info]),
     {next_state, StateName, StateData}.
 
-%%----------------------------------------------------------------------
-%% Func: terminate/3
-%% Purpose: Shutdown the fsm
-%% Returns: any
-%%----------------------------------------------------------------------
 terminate(Reason, StateName, StateData) ->
     ?INFO_MSG("terminated: ~p", [Reason]),
     case StateName of
@@ -387,6 +283,7 @@ terminate(Reason, StateName, StateData) ->
 			dict:fetch_keys(StateData#state.host_opts));
       _ -> ok
     end,
+    catch send_trailer(StateData),
     (StateData#state.sockmod):close(StateData#state.socket),
     ok.
 
@@ -401,13 +298,69 @@ print_state(State) -> State.
 %%% Internal functions
 %%%----------------------------------------------------------------------
 
+-spec send_text(state(), iodata()) -> ok.
 send_text(StateData, Text) ->
     (StateData#state.sockmod):send(StateData#state.socket,
 				   Text).
 
+-spec send_element(state(), xmpp_element()) -> ok.
 send_element(StateData, El) ->
-    send_text(StateData, fxml:element_to_binary(El)).
+    El1 = fix_ns(xmpp:encode(El)),
+    send_text(StateData, fxml:element_to_binary(El1)).
+
+-spec send_error(state(), xmlel() | stanza(), error()) -> ok.
+send_error(StateData, Stanza, Error) ->
+    Type = xmpp:get_type(Stanza),
+    if Type == error; Type == result;
+       Type == <<"error">>; Type == <<"result">> ->
+	    ok;
+       true ->
+	    send_element(StateData, xmpp:make_error(Stanza, Error))
+    end.
+
+-spec send_header(state(), binary()) -> ok.
+send_header(StateData, Host) ->
+    send_text(StateData,
+	      io_lib:format(
+		<<"<?xml version='1.0'?><stream:stream "
+		  "xmlns:stream='http://etherx.jabber.org/stream"
+		  "s' xmlns='jabber:component:accept' id='~s' "
+		  "from='~s'>">>,
+		[StateData#state.streamid, fxml:crypt(Host)])).
+
+-spec send_trailer(state()) -> ok.
+send_trailer(StateData) ->
+    send_text(StateData, <<"</stream:stream>">>).
+
+-spec fix_ns(xmlel()) -> xmlel().
+fix_ns(#xmlel{name = Name} = El) when Name == <<"message">>;
+                                      Name == <<"iq">>;
+                                      Name == <<"presence">> ->
+    Attrs = lists:filter(
+              fun({<<"xmlns">>, _}) -> false;
+                 (_) -> true
+              end, El#xmlel.attrs),
+    El#xmlel{attrs = Attrs};
+fix_ns(El) ->
+    El.
+
+-spec decode_element(xmlel(), state_name(), state()) -> fsm_transition().
+decode_element(#xmlel{} = El, StateName, StateData) ->
+    try xmpp:decode(El, [ignore_els]) of
+	Pkt -> ?MODULE:StateName(Pkt, StateData)
+    catch error:{xmpp_codec, Why} ->
+            case xmpp:is_stanza(El) of
+                true ->
+                    Lang = xmpp:get_lang(El),
+                    Txt = xmpp:format_error(Why),
+                    send_error(StateData, El, xmpp:err_bad_request(Txt, Lang));
+                false ->
+                    ok
+            end,
+            {next_state, StateName, StateData}
+    end.
 
+-spec new_id() -> binary().
 new_id() -> randoms:get_string().
 
 transform_listen_option({hosts, Hosts, O}, Opts) ->