Implement SM resumption into SASL2

big_tests/tests/sasl2_SUITE.erl

@@ -22,7 +22,8 @@ groups() ->
      {all_tests, [parallel],
       [
        {group, basic},
-       {group, scram}
+       {group, scram},
+       {group, stream_management}
       ]},
      {basic, [parallel],
       [
@@ -40,6 +41,10 @@ groups() ->
        authenticate_with_scram_bad_abort,
        authenticate_with_scram_bad_response,
        authenticate_with_scram
+      ]},
+     {stream_management, [parallel],
+      [
+       stream_resumption_is_bound_at_sasl2_success
       ]}
     ].
 
@@ -78,7 +83,9 @@ end_per_testcase(Name, Config) ->
 load_sasl_extensible(Config) ->
     HostType = domain_helper:host_type(),
     Config1 = dynamic_modules:save_modules(HostType, Config),
-    dynamic_modules:ensure_modules(HostType, [{mod_sasl2, #{}}]),
+    Modules = [{mod_sasl2, config_parser_helper:default_mod_config(mod_sasl2)},
+               {mod_stream_management, config_parser_helper:mod_config(mod_stream_management, #{ack_freq => never})}],
+    dynamic_modules:ensure_modules(HostType, Modules),
     Config1.
 
 %%--------------------------------------------------------------------
@@ -159,3 +166,11 @@ authenticate_again_results_in_stream_error(Config) ->
              plain_authentication, receive_features, plain_authentication],
     #{answer := Response} = sasl2_helper:apply_steps(Steps, Config),
     escalus:assert(is_stream_error, [<<"policy-violation">>, <<>>], Response).
+
+stream_resumption_is_bound_at_sasl2_success(Config) ->
+    Steps = [buffer_messages_and_die, connect_tls_user, start_stream_get_features,
+             auth_with_resumption],
+    #{answer := Success, smid := SMID, smh := _Tag} = sasl2_helper:apply_steps(Steps, Config),
+    ?assertMatch(#xmlel{name = <<"success">>, attrs = [{<<"xmlns">>, ?NS_SASL_2}]}, Success),
+    Resumed = exml_query:path(Success, [{element_with_ns, <<"resumed">>, ?NS_STREAM_MGNT_3}]),
+    ?assert(escalus_pred:is_sm_resumed(SMID, Resumed)).

big_tests/tests/sasl2_helper.erl

@@ -58,6 +58,12 @@ send_bad_user_agent(_Config, Client, Data) ->
     Answer = escalus_client:wait_for_stanza(Client),
     {Client, Data#{answer => Answer}}.
 
+auth_with_resumption(Config, Client, #{smid := SMID, texts := Texts} = Data) ->
+    Resume = escalus_stanza:resume(SMID, 1),
+    {Client, Data} = plain_auth(Config, Client, Data, [Resume]),
+    Msgs = sm_helper:wait_for_messages(Client, Texts),
+    Data#{sm_storage => Msgs}.
+
 plain_auth_user_agent_without_id(Config, Client, Data) ->
     plain_auth(Config, Client, Data, [user_agent_elem_without_id()]).
 
@@ -124,6 +130,22 @@ scram_step_2(_Config, Client,
         {error, _, _} -> throw({auth_failed, SuccessStanza})
     end.
 
+buffer_messages_and_die(Config, _Client, Data) ->
+    Spec = escalus_fresh:create_fresh_user(Config, alice),
+    Client = sm_helper:connect_spec(Spec, sr_presence, manual),
+    C2SPid = mongoose_helper:get_session_pid(Client),
+    BobSpec = escalus_fresh:create_fresh_user(Config, bob),
+    {ok, Bob, _} = escalus_connection:start(BobSpec),
+    Texts = [ integer_to_binary(N) || N <- [1, 2, 3]],
+    sm_helper:send_messages(Bob, Client, Texts),
+    %% Client receives them, but doesn't ack.
+    sm_helper:wait_for_messages(Client, Texts),
+    %% Client's connection is violently terminated.
+    escalus_client:kill_connection(Config, Client),
+    sm_SUITE:wait_until_resume_session(C2SPid),
+    SMID = sm_helper:client_to_smid(Client),
+    {C2SPid, Data#{spec => Spec, smid => SMID, smh => 3, texts => Texts}}.
+
 receive_features(_Config, Client, Data) ->
     Features = escalus_client:wait_for_stanza(Client),
     {Client, Data#{features => Features}}.

src/c2s/mongoose_c2s_acc.erl

@@ -13,6 +13,7 @@
 %% - `route': mongoose_acc elements to trigger the whole `handle_route' pipeline.
 %% - `flush': mongoose_acc elements to trigger the `handle_flush` pipeline.
 %% - `socket_send': xml elements to send on the socket to the user.
+%% - `socket_send_first': xml elements to send on the socket to the user, but appended first.
 -module(mongoose_c2s_acc).
 
 -export([new/0, new/1,
@@ -21,7 +22,16 @@
          to_acc/3, to_acc_many/2
         ]).
 
--type key() :: state_mod | actions | c2s_state | c2s_data | stop | hard_stop | route | flush | socket_send.
+-type key() :: state_mod
+             | actions
+             | c2s_state
+             | c2s_data
+             | stop
+             | hard_stop
+             | route
+             | flush
+             | socket_send
+             | socket_send_first.
 -type pairs() :: [pair()].
 -type pair() :: {state_mod, {module(), term()}}
               | {actions, gen_statem:action()}
@@ -31,7 +41,8 @@
               | {hard_stop, term() | {shutdown, atom()}}
               | {route, mongoose_acc:t()}
               | {flush, mongoose_acc:t()}
-              | {socket_send, exml:element()}.
+              | {socket_send, exml:element() | [exml:element()]}
+              | {socket_send_first, exml:element() | [exml:element()]}.
 
 -type t() :: #{
         state_mod := #{module() => term()},
@@ -122,7 +133,8 @@ from_mongoose_acc(Acc, Key) ->
             (mongoose_acc:t(), stop, atom() | {shutdown, atom()}) -> mongoose_acc:t();
             (mongoose_acc:t(), route, mongoose_acc:t()) -> mongoose_acc:t();
             (mongoose_acc:t(), flush, mongoose_acc:t()) -> mongoose_acc:t();
-            (mongoose_acc:t(), socket_send, exml:element()) -> mongoose_acc:t().
+            (mongoose_acc:t(), socket_send, exml:element() | [exml:element()]) -> mongoose_acc:t();
+            (mongoose_acc:t(), socket_send_first, exml:element() | [exml:element()]) -> mongoose_acc:t().
 to_acc(Acc, Key, NewValue) ->
     C2SAcc = mongoose_acc:get_statem_acc(Acc),
     C2SAcc1 = to_c2s_acc(C2SAcc, {Key, NewValue}),
@@ -161,6 +173,10 @@ to_c2s_acc(C2SAcc = #{socket_send := Stanzas}, {socket_send, NewStanzas}) when i
     C2SAcc#{socket_send := lists:reverse(NewStanzas) ++ Stanzas};
 to_c2s_acc(C2SAcc = #{socket_send := Stanzas}, {socket_send, Stanza}) ->
     C2SAcc#{socket_send := [Stanza | Stanzas]};
+to_c2s_acc(C2SAcc = #{socket_send := Stanzas}, {socket_send_first, NewStanzas}) when is_list(NewStanzas) ->
+    C2SAcc#{socket_send := Stanzas ++ NewStanzas};
+to_c2s_acc(C2SAcc = #{socket_send := Stanzas}, {socket_send_first, Stanza}) ->
+    C2SAcc#{socket_send := Stanzas ++ [Stanza]};
 to_c2s_acc(C2SAcc = #{actions := Actions}, {stop, Reason}) ->
     C2SAcc#{actions := [{next_event, cast, {stop, Reason}} | Actions]};
 to_c2s_acc(C2SAcc, {Key, NewValue}) ->

src/mod_sasl2.erl

@@ -25,14 +25,20 @@
          sasl2_failure/3,
          sasl2_error/3]).
 
+%% helpers
+-export([get_inline_request/2, put_inline_request/3, update_inline_request/4]).
+
 -type maybe_binary() :: undefined | binary().
+-type status() :: pending | success | failure.
+-type inline_request() :: #{request := exml:element(),
+                            response := undefined | exml:element(),
+                            status := status()}.
 -type mod_state() :: #{authenticated := boolean(),
                        id := not_provided | uuid:uuid(),
                        software := not_provided | binary(),
                        device := not_provided | binary()}.
 
--type params() :: #{c2s_data => mongoose_c2s:data(), c2s_state => mongoose_c2s:state()}.
--export_type([params/0]).
+-export_type([inline_request/0]).
 
 %% gen_mod
 -spec start(mongooseim:host_type(), gen_mod:module_opts()) -> ok.
@@ -241,27 +247,72 @@ handle_sasl_step(HookParams, {error, NewSaslAcc, Result}, Retries) ->
 -spec handle_sasl_success(
         mongoose_c2s_hooks:params(), mongoose_acc:t(), mongoose_c2s_sasl:success()) ->
     mongoose_c2s:fsm_res().
-handle_sasl_success(#{c2s_data := C2SData, c2s_state := C2SState} = HookParams, SaslAcc,
-                    #{server_out := MaybeServerOut, jid := Jid, auth_module := AuthMod}) ->
-    C2SData1 = mongoose_c2s:set_jid(C2SData, Jid),
-    C2SData2 = mongoose_c2s:set_auth_module(C2SData1, AuthMod),
-    HostType = mongoose_c2s:get_host_type(C2SData2),
+handle_sasl_success(#{c2s_data := C2SData} = HookParams,
+                    SaslAcc, #{server_out := MaybeServerOut, jid := Jid, auth_module := AuthMod}) ->
+    C2SData1 = build_final_c2s_data(C2SData, Jid, AuthMod),
     ?LOG_INFO(#{what => auth_success, text => <<"Accepted SASL authentication">>,
-                user => jid:to_binary(Jid), c2s_state => C2SData2}),
-    ModState = get_mod_state(C2SData2),
-    Actions = [pop_callback_module,
-               {next_event, internal, mongoose_c2s_stanzas:stream_header(C2SData2)},
-               mongoose_c2s:state_timeout(C2SData2)],
-    SuccessStanza = success_stanza(Jid, MaybeServerOut),
-    StreamFeaturesStanza = mongoose_c2s_stanzas:stream_features_after_auth(C2SData2),
+                user => jid:to_binary(Jid), c2s_state => C2SData1}),
     C2SState1 = {wait_for_feature_after_auth, ?BIND_RETRIES},
-    ToAcc = [{socket_send, [SuccessStanza, StreamFeaturesStanza]},
-             {actions, Actions},
-             {state_mod, {?MODULE, ModState#{authenticated := true}}},
-             {c2s_state, C2SState1}],
-    SaslAcc1 = mongoose_c2s_acc:to_acc_many(SaslAcc, ToAcc),
-    SaslAcc2 = mongoose_hooks:sasl2_success(HostType, SaslAcc1, HookParams),
-    mongoose_c2s:handle_state_after_packet(C2SData2, C2SState, SaslAcc2).
+    SaslAcc1 = extend_sasl2_acc(SaslAcc, C2SData),
+    SaslAcc2 = run_sasl2_success(SaslAcc1, C2SData1, HookParams),
+    SaslAcc3 = process_sasl2_success_result(SaslAcc2, C2SData1, Jid, MaybeServerOut),
+    mongoose_c2s:handle_state_after_packet(C2SData1, C2SState1, SaslAcc3).
+
+-spec extend_sasl2_acc(mongoose_acc:t(), mongoose_c2s:data()) -> mongoose_acc:t().
+extend_sasl2_acc(SaslAcc, C2SData) ->
+    ModState = get_mod_state(C2SData),
+    Actions = [pop_callback_module,
+               {next_event, internal, mongoose_c2s_stanzas:stream_header(C2SData)},
+               mongoose_c2s:state_timeout(C2SData)],
+    ToAcc = [{actions, Actions},
+             {state_mod, {?MODULE, ModState#{authenticated := true}}}],
+    mongoose_c2s_acc:to_acc_many(SaslAcc, ToAcc).
+
+-spec run_sasl2_success(mongoose_acc:t(), mongoose_c2s:data(), mongoose_c2s_hooks:params()) ->
+    mongoose_acc:t().
+run_sasl2_success(SaslAcc, C2SData, HookParams) ->
+    HostType = mongoose_c2s:get_host_type(C2SData),
+    mongoose_hooks:sasl2_success(HostType, SaslAcc, HookParams).
+
+-spec process_sasl2_success_result(mongoose_acc:t(), mongoose_c2s:data(), jid:jid(), maybe_binary()) ->
+    mongoose_acc:t().
+process_sasl2_success_result(SaslAcc, C2SData, Jid, MaybeServerOut) ->
+    Inlines = get_inline_requests(SaslAcc),
+    InlineAnswers = get_inline_responses(Inlines),
+    SuccessStanza = success_stanza(Jid, MaybeServerOut, InlineAnswers),
+    ToAcc = case is_session_established(SaslAcc) of
+                true ->
+                    [{socket_send_first, SuccessStanza}];
+                false ->
+                    StreamFeaturesStanza = mongoose_c2s_stanzas:stream_features_after_auth(C2SData),
+                    [{socket_send_first, SuccessStanza}, {socket_send, StreamFeaturesStanza}]
+            end,
+    mongoose_c2s_acc:to_acc_many(SaslAcc, ToAcc).
+
+-spec is_session_established(mongoose_acc:t()) -> boolean().
+is_session_established(SaslAcc) ->
+    #{c2s_state := State} = mongoose_c2s_acc:get_statem_result(SaslAcc),
+    session_established =:= State.
+
+-spec get_inline_responses([inline_request()]) -> [exml:element()].
+get_inline_responses(Inlines) ->
+    [ Response || {_Module, #{status := Status, response := Response}} <- Inlines,
+                  pending =/= Status,
+                  undefined =/= Response ].
+
+-spec success_stanza(jid:jid(), maybe_binary(), [exml:element()]) -> exml:element().
+success_stanza(AuthJid, undefined, Inlines) ->
+    AuthorizationId = success_subelement(<<"authorization-identifier">>, jid:to_binary(AuthJid)),
+    sasl2_ns_stanza(<<"success">>, [AuthorizationId | Inlines]);
+success_stanza(AuthJid, CData, Inlines) ->
+    AdditionalData = success_subelement(<<"additional-data">>, base64:encode(CData)),
+    AuthorizationId = success_subelement(<<"authorization-identifier">>, jid:to_binary(AuthJid)),
+    sasl2_ns_stanza(<<"success">>, [AdditionalData, AuthorizationId | Inlines]).
+
+-spec build_final_c2s_data(mongoose_c2s:data(), jid:jid(), module()) -> mongoose_c2s:data().
+build_final_c2s_data(C2SData, Jid, AuthMod) ->
+    C2SData1 = mongoose_c2s:set_jid(C2SData, Jid),
+    mongoose_c2s:set_auth_module(C2SData1, AuthMod).
 
 -spec handle_sasl_continue(
         mongoose_c2s_hooks:params(), mongoose_acc:t(), mongoose_c2s_sasl:continue(), mongoose_c2s:retries()) ->
@@ -308,15 +359,6 @@ handle_sasl_error(#{c2s_data := C2SData} = HookParams, SaslAcc, #{type := Type,
     SaslAcc1 = mongoose_hooks:sasl2_error(HostType, SaslAcc, HookParams),
     {stop, mongoose_c2s_acc:to_acc(SaslAcc1, hard_stop, sasl2_violation)}.
 
--spec success_stanza(jid:jid(), maybe_binary()) -> exml:element().
-success_stanza(AuthJid, undefined) ->
-    AuthorizationId = success_subelement(<<"authorization-identifier">>, jid:to_binary(AuthJid)),
-    sasl2_ns_stanza(<<"success">>, [AuthorizationId]);
-success_stanza(AuthJid, CData) ->
-    AdditionalData = success_subelement(<<"additional-data">>, base64:encode(CData)),
-    AuthorizationId = success_subelement(<<"authorization-identifier">>, jid:to_binary(AuthJid)),
-    sasl2_ns_stanza(<<"success">>, [AdditionalData, AuthorizationId]).
-
 -spec challenge_stanza(binary()) -> exml:element().
 challenge_stanza(ServerOut) ->
     Challenge = #xmlcdata{content = base64:encode(ServerOut)},
@@ -390,3 +432,26 @@ inlines(InlineFeatures) ->
 -spec feature_name() -> binary().
 feature_name() ->
     <<"authentication">>.
+
+-spec get_inline_requests(mongoose_acc:t()) -> [{module(), inline_request()}].
+get_inline_requests(SaslAcc) ->
+    mongoose_acc:get(?MODULE, SaslAcc).
+
+-spec get_inline_request(mongoose_acc:t(), module()) -> undefined | inline_request().
+get_inline_request(SaslAcc, ModuleRequest) ->
+    mongoose_acc:get(?MODULE, ModuleRequest, undefined, SaslAcc).
+
+-spec put_inline_request(mongoose_acc:t(), module(), exml:element()) -> mongoose_acc:t().
+put_inline_request(SaslAcc, ModuleRequest, XmlRequest) ->
+    Request = #{request => XmlRequest, response => undefined, status => pending},
+    mongoose_acc:set(?MODULE, ModuleRequest, Request, SaslAcc).
+
+-spec update_inline_request(mongoose_acc:t(), module(), exml:element(), status()) -> mongoose_acc:t().
+update_inline_request(SaslAcc, ModuleRequest, XmlResponse, Status) ->
+    case mongoose_acc:get(?MODULE, ModuleRequest, undefined, SaslAcc) of
+        undefined ->
+            SaslAcc;
+        Request ->
+            Request1 = Request#{response := XmlResponse, status := Status},
+            mongoose_acc:set(?MODULE, ModuleRequest, Request1, SaslAcc)
+    end.