Add options to have separate parser limits before authentication

This also add new option that allows to limit max number of
elements that stanza can have.

Author: prefiks

rebar.config

@@ -37,7 +37,7 @@
 	      {"priv/lib/xmpp_lang.so", ["c_src/xmpp_lang.c"]}]}.
 
 {deps, [{p1_utils, "~> 1.0.29", {git, "https://github.com/processone/p1_utils", {tag, "1.0.29"}}},
-	{fast_xml, "~> 1.1.58", {git, "https://github.com/processone/fast_xml", {tag, "1.1.58"}}},
+	{fast_xml, "~> 1.1.59", {git, "https://github.com/processone/fast_xml", {tag, "1.1.59"}}},
 	{fast_tls, "~> 1.1.26", {git, "https://github.com/processone/fast_tls", {tag, "1.1.26"}}},
 	{ezlib, "~> 1.0.16", {git, "https://github.com/processone/ezlib", {tag, "1.0.16"}}},
 	{idna, "~> 7.1", {git, "https://github.com/benoitc/erlang-idna", {tag, "7.1.0"}}},

src/xmpp_socket.erl

@@ -51,7 +51,8 @@
 	 release/1,
 	 get_tls_cb_exporter/1,
 	 get_tls_cert_hash/1,
-	 finish_tls_handshake/1]).
+	 finish_tls_handshake/1,
+	 change_limits/3]).
 
 -include("xmpp.hrl").
 -include_lib("public_key/include/public_key.hrl").
@@ -71,6 +72,7 @@
 -record(socket_state, {sockmod           :: sockmod(),
                        socket            :: socket(),
 		       max_stanza_size   :: timeout(),
+		       max_stanza_elements :: timeout(),
 		       xml_stream :: undefined | fxml_stream:xml_stream_state(),
 		       shaper = none :: none | p1_shaper:state(),
 		       sock_peer_name = none :: none | {endpoint(), endpoint()},
@@ -108,17 +110,19 @@
 -spec new(sockmod(), socket(), [proplists:property()]) -> socket_state().
 new(SockMod, Socket, Opts) ->
     MaxStanzaSize = proplists:get_value(max_stanza_size, Opts, infinity),
+    MaxStanzaElements = proplists:get_value(max_stanza_elements, Opts, infinity),
     SockPeer =  proplists:get_value(sock_peer_name, Opts, none),
     XMLStream = case get_owner(SockMod, Socket) of
 		    Pid when Pid == self() ->
-			fxml_stream:new(self(), MaxStanzaSize);
+			fxml_stream:new(self(), {MaxStanzaSize, MaxStanzaElements});
 		    _ ->
 			undefined
 		end,
     #socket_state{sockmod = SockMod,
 		  socket = Socket,
 		  xml_stream = XMLStream,
 		  max_stanza_size = MaxStanzaSize,
+		  max_stanza_elements = MaxStanzaElements,
 		  sock_peer_name = SockPeer}.
 
 connect(Addr, Port, Opts) ->
@@ -227,19 +231,28 @@ compress(_, _) ->
 
 reset_stream(#socket_state{xml_stream = XMLStream,
 			   sockmod = SockMod, socket = Socket,
-			   max_stanza_size = MaxStanzaSize} = SocketData) ->
+			   max_stanza_size = MaxStanzaSize,
+			   max_stanza_elements = MaxStanzaElements} = SocketData) ->
     if XMLStream /= undefined ->
 	    XMLStream1 = try fxml_stream:reset(XMLStream)
 			 catch error:_ ->
 				 fxml_stream:close(XMLStream),
-				 fxml_stream:new(self(), MaxStanzaSize)
+				 fxml_stream:new(self(), {MaxStanzaSize, MaxStanzaElements})
 			 end,
 	    SocketData#socket_state{xml_stream = XMLStream1};
        true ->
 	    Socket1 = SockMod:reset_stream(Socket),
 	    SocketData#socket_state{socket = Socket1}
     end.
 
+change_limits(#socket_state{xml_stream = undefined} = SocketData, _NewMaxStanzaSize, _NewMaxStanzaElements) ->
+	SocketData;
+change_limits(#socket_state{xml_stream = XMLStream} = SocketData, NewMaxStanzaSize, NewMaxStanzaElements) ->
+    fxml_stream:change_limits(XMLStream, NewMaxStanzaSize, NewMaxStanzaElements),
+	SocketData#socket_state{max_stanza_size = NewMaxStanzaSize, max_stanza_elements = NewMaxStanzaElements};
+change_limits(SocketData, _, _) ->
+	SocketData.
+
 -spec send_element(socket_state(), fxml:xmlel()) -> ok | {error, inet:posix()}.
 send_element(#socket_state{xml_stream = undefined} = SocketData, El) ->
     send_xml(SocketData, {xmlstreamelement, El});

src/xmpp_stream_in.erl

@@ -302,7 +302,12 @@ init([Mod, {SockMod, Socket}, Opts]) ->
 handle_cast(accept, #{socket := Socket,
 		      socket_mod := SockMod,
 		      socket_opts := Opts} = State) ->
-    XMPPSocket = xmpp_socket:new(SockMod, Socket, Opts),
+    MaxStanzaSize = proplists:get_value(max_stanza_size, Opts, infinity),
+    MaxStanzaElements = proplists:get_value(max_stanza_elements, Opts, infinity),
+    PreAuthMaxStanzaSize = proplists:get_value(pre_auth_max_stanza_size, Opts, MaxStanzaSize),
+	PreAuthMaxStanzaElements = proplists:get_value(pre_auth_max_stanza_elements, Opts, MaxStanzaElements),
+	Opts2 = [{max_stanza_size, PreAuthMaxStanzaSize}, {max_stanza_elements, PreAuthMaxStanzaElements} | Opts],
+    XMPPSocket = xmpp_socket:new(SockMod, Socket, Opts2),
     SocketMonitor = xmpp_socket:monitor(XMPPSocket),
     case xmpp_socket:peername(XMPPSocket) of
 	{ok, IP} ->
@@ -314,9 +319,14 @@ handle_cast(accept, #{socket := Socket,
 		{stop, State4} ->
 		    {stop, normal, State4};
 		State4 ->
-		    case is_disconnected(State4) of
-			true -> noreply(State4);
-			false -> handle_info({tcp, Socket, <<>>}, State4)
+			State5 = if MaxStanzaSize /= PreAuthMaxStanzaSize orelse MaxStanzaElements /= PreAuthMaxStanzaElements ->
+					State4#{post_auth_stanza_limits => {MaxStanzaSize, MaxStanzaElements}};
+				true ->
+					State4
+			end,
+		    case is_disconnected(State5) of
+			true -> noreply(State5);
+			false -> handle_info({tcp, Socket, <<>>}, State5)
 		    end
 	    end;
 	{error, _} ->
@@ -862,10 +872,17 @@ process_stream_established(State) ->
     State1 = State#{stream_authenticated => true,
 		    stream_state => established,
 		    stream_timeout => infinity},
-    try callback(handle_stream_established, State1)
-    catch _:{?MODULE, undef} -> State1
+    State2 = update_parser_limits(State1),
+    try callback(handle_stream_established, State2)
+    catch _:{?MODULE, undef} -> State2
     end.
 
+update_parser_limits(#{socket := Socket, post_auth_stanza_limits := {MaxStanzaSize, MaxElementsCount}} = State) ->
+    Socket2 = xmpp_socket:change_limits(Socket, MaxStanzaSize, MaxElementsCount),
+	maps:remove(post_auth_stanza_limits, State#{socket => Socket2});
+update_parser_limits(State) ->
+    State.
+
 -spec process_compress(compress(), state()) -> state().
 process_compress(#compress{},
 		 #{stream_compressed := Compressed,
@@ -1032,8 +1049,9 @@ process_sasl_success(Props, ServerOut,
 	    case is_disconnected(State2) of
 		true -> State2;
 		false ->
-		    State3 = map_remove_keys(State2, [sasl_state, sasl_mech, sasl_channel_bindings]),
-		    State3#{stream_id => xmpp_stream:new_id(),
+		    State3 = update_parser_limits(State2),
+		    State4 = map_remove_keys(State3, [sasl_state, sasl_mech, sasl_channel_bindings]),
+		    State4#{stream_id => xmpp_stream:new_id(),
 			    stream_authenticated => true,
 			    stream_header_sent => false,
 			    stream_restarted => true,
@@ -1161,13 +1179,14 @@ process_sasl2_success(Props, ServerOut,
 	    case is_disconnected(State1) of
 		true -> State1;
 		false ->
-		    State2 = State1#{stream_id => xmpp_stream:new_id(),
-				     stream_authenticated => true,
-				     user => User},
+		    State2 = update_parser_limits(
+				State1#{stream_id => xmpp_stream:new_id(),
+					stream_authenticated => true,
+					user => User}),
 		    {State3, NewEls, Results} =
-		    try callback(handle_sasl2_inline, InlineEls, State2)
-		    catch _:{?MODULE, undef} -> {State2, InlineEls, []}
-		    end,
+			try callback(handle_sasl2_inline, InlineEls, State2)
+			catch _:{?MODULE, undef} -> {State2, InlineEls, []}
+			end,
 
 		    case NewEls of
 			{continue, Tasks} ->