Merge branch 'anders/diameter/tbit/OTP-10619'

* anders/diameter/tbit/OTP-10619:
  Add discard testcases to failover suite
  Check for T flag in failover suite
  Fix setting of T flag at failover

2 files changed, 99 insertions, 52 deletions

diff --git a/lib/diameter/src/base/diameter_service.erl b/lib/diameter/src/base/diameter_service.erl
index a4a0b80348..91384b8b91 100644
--- a/lib/diameter/src/base/diameter_service.erl
+++ b/lib/diameter/src/base/diameter_service.erl
@@ -1703,16 +1703,20 @@ send(Pid, Pkt) ->
 %% retransmit/4
 
 retransmit({TPid, Caps, #diameter_app{alias = Alias} = App} = T,
-           #request{app = Alias, packet = Pkt}
+           #request{app = Alias, packet = Pkt0}
            = Req,
            SvcName,
            Timeout) ->
-    have_request(Pkt, TPid)      %% Don't failover to a peer we've
+    have_request(Pkt0, TPid)      %% Don't failover to a peer we've
         andalso ?THROW(timeout), %% already sent to.
 
+    #diameter_packet{header = Hdr0} = Pkt0,
+    Hdr = Hdr0#diameter_header{is_retransmitted = true},
+    Pkt = Pkt0#diameter_packet{header = Hdr},
+
     resend_req(cb(App, prepare_retransmit, [Pkt, SvcName, {TPid, Caps}]),
                T,
-               Req,
+               Req#request{packet = Pkt},
                Timeout,
                []).
 
diff --git a/lib/diameter/test/diameter_failover_SUITE.erl b/lib/diameter/test/diameter_failover_SUITE.erl
index ed31670031..bb820a8bf2 100644
--- a/lib/diameter/test/diameter_failover_SUITE.erl
+++ b/lib/diameter/test/diameter_failover_SUITE.erl
@@ -18,19 +18,19 @@
 %%
 
 %%
-%% Tests of traffic between six Diameter nodes in three realms,
+%% Tests of traffic between seven Diameter nodes in four realms,
 %% connected as follows.
 %%
-%%                  ----- SERVER1.REALM2
-%%                /
-%%               /  ----- SERVER2.REALM2
-%%              | /
-%%   CLIENT.REALM1 ------ SERVER3.REALM2
-%%              | \
-%%              |  \
-%%               \   ---- SERVER1.REALM3
-%%                \
-%%                  ----- SERVER2.REALM3
+%%                  ----- SERVER1.REALM2 -----
+%%                /                            \
+%%               /  ----- SERVER2.REALM2 -----  \
+%%              | /                            \ |
+%%   CLIENT.REALM1 ------ SERVER3.REALM2 ------ CLIENT.REALM4
+%%              | \                            / |
+%%              |  \                          /  |
+%%               \   ---- SERVER1.REALM3 -----  /
+%%                \                            /
+%%                  ----- SERVER2.REALM3 -----
 %%
 
 -module(diameter_failover_SUITE).
@@ -44,12 +44,16 @@
          connect/1,
          send_ok/1,
          send_nok/1,
+         send_discard_1/1,
+         send_discard_2/1,
          stop_services/1,
          stop/1]).
 
 %% diameter callbacks
 -export([pick_peer/4,
          prepare_request/3,
+         prepare_retransmit/3,
+         handle_error/4,
          handle_answer/4,
          handle_request/3]).
 
@@ -62,14 +66,18 @@
 
 -define(ADDR, {127,0,0,1}).
 
--define(CLIENT,  "CLIENT.REALM1").
+-define(CLIENT1, "CLIENT.REALM1").
+-define(CLIENT2, "CLIENT.REALM4").
 -define(SERVER1, "SERVER1.REALM2").
 -define(SERVER2, "SERVER2.REALM2").
 -define(SERVER3, "SERVER3.REALM2").
 -define(SERVER4, "SERVER1.REALM3").
 -define(SERVER5, "SERVER2.REALM3").
 
--define(SERVICES, [?CLIENT, ?SERVER1, ?SERVER2, ?SERVER3, ?SERVER4, ?SERVER5]).
+-define(IS_CLIENT(Svc), Svc == ?CLIENT1; Svc == ?CLIENT2).
+
+-define(CLIENTS, [?CLIENT1, ?CLIENT2]).
+-define(SERVERS, [?SERVER1, ?SERVER2, ?SERVER3, ?SERVER4, ?SERVER5]).
 
 -define(DICT_COMMON,  ?DIAMETER_DICT_COMMON).
 
@@ -77,26 +85,27 @@
 -define(APP_ID, ?DICT_COMMON:id()).
 
 %% Config for diameter:start_service/2.
--define(SERVICE(Host, Dict),
+-define(SERVICE(Host),
         [{'Origin-Host', Host},
          {'Origin-Realm', realm(Host)},
          {'Host-IP-Address', [?ADDR]},
          {'Vendor-Id', 12345},
          {'Product-Name', "OTP/diameter"},
-         {'Acct-Application-Id', [Dict:id()]},
+         {'Acct-Application-Id', [?APP_ID]},
          {application, [{alias, ?APP_ALIAS},
-                        {dictionary, Dict},
+                        {dictionary, ?DICT_COMMON},
                         {module, #diameter_callback
                                   {peer_up = false,
                                    peer_down = false,
-                                   handle_error = false,
-                                   prepare_retransmit = false,
                                    default = ?MODULE}},
                         {answer_errors, callback}]}]).
 
 -define(SUCCESS, 2001).
 
--define(LOGOUT, ?'DIAMETER_BASE_TERMINATION-CAUSE_DIAMETER_LOGOUT').
+%% Value of Termination-Cause determines client/server behaviour.
+-define(LOGOUT,   ?'DIAMETER_BASE_TERMINATION-CAUSE_DIAMETER_LOGOUT').
+-define(MOVED,    ?'DIAMETER_BASE_TERMINATION-CAUSE_DIAMETER_USER_MOVED').
+-define(TIMEOUT,  ?'DIAMETER_BASE_TERMINATION-CAUSE_DIAMETER_SESSION_TIMEOUT').
 
 %% ===========================================================================
 
@@ -109,6 +118,8 @@ all() ->
      connect,
      send_ok,
      send_nok,
+     send_discard_1,
+     send_discard_2,
      stop_services,
      stop].
 
@@ -119,19 +130,20 @@ start(_Config) ->
     ok = diameter:start().
 
 start_services(_Config) ->
-    S = [server(N, ?DICT_COMMON) || N <- tl(?SERVICES)],
-
-    ok = diameter:start_service(?CLIENT, ?SERVICE(?CLIENT, ?DICT_COMMON)),
+    Servers = [server(N) || N <- ?SERVERS],
+    [] = [T || C <- ?CLIENTS,
+               T <- [diameter:start_service(C, ?SERVICE(C))],
+               T /= ok],
 
-    {save_config, [{?CLIENT, S}]}.
+    {save_config, Servers}.
 
 connect(Config) ->
-    {_, Conns} = proplists:get_value(saved_config, Config),
+    {start_services, Servers} = proplists:get_value(saved_config, Config),
 
-    lists:foreach(fun({CN,Ss}) -> connect(CN, Ss) end, Conns).
+    lists:foreach(fun(C) -> connect(C, Servers) end, ?CLIENTS).
 
 stop_services(_Config) ->
-    [] = [{H,T} || H <- ?SERVICES,
+    [] = [{H,T} || H <- ?CLIENTS ++ ?SERVERS,
                    T <- [diameter:stop_service(H)],
                    T /= ok].
 
@@ -140,8 +152,8 @@ stop(_Config) ->
 
 %% ----------------------------------------
 
-server(Name, Dict) ->
-    ok = diameter:start_service(Name, ?SERVICE(Name, Dict)),
+server(Name) ->
+    ok = diameter:start_service(Name, ?SERVICE(Name)),
     {Name, ?util:listen(Name, tcp)}.
 
 connect(Name, Refs) ->
@@ -153,30 +165,39 @@ connect(Name, Refs) ->
 %% Send an STR and expect success after SERVER3 answers after a couple
 %% of failovers.
 send_ok(_Config) ->
-    Req = ['STR', {'Destination-Realm', realm(?SERVER1)},
-                  {'Termination-Cause', ?LOGOUT},
-                  {'Auth-Application-Id', ?APP_ID}],
+    Req = #diameter_base_STR{'Destination-Realm' = realm(?SERVER1),
+                             'Termination-Cause' = ?LOGOUT,
+                             'Auth-Application-Id' = ?APP_ID},
     #diameter_base_STA{'Result-Code' = ?SUCCESS,
                        'Origin-Host' = ?SERVER3}
-        = call(Req, [{filter, realm}]).
+        = call(?CLIENT1, Req).
 
 %% Send an STR and expect failure when both servers fail.
 send_nok(_Config) ->
-    Req = ['STR', {'Destination-Realm', realm(?SERVER4)},
-                  {'Termination-Cause', ?LOGOUT},
-                  {'Auth-Application-Id', ?APP_ID}],
-    {error, failover} = call(Req, [{filter, realm}]).
+    Req = #diameter_base_STR{'Destination-Realm' = realm(?SERVER4),
+                             'Termination-Cause' = ?LOGOUT,
+                             'Auth-Application-Id' = ?APP_ID},
+    {failover, ?LOGOUT} = call(?CLIENT1, Req).
+
+%% Send an STR and have prepare_retransmit discard it.
+send_discard_1(_Config) ->
+    Req = #diameter_base_STR{'Destination-Realm' = realm(?SERVER1),
+                             'Termination-Cause' = ?TIMEOUT,
+                             'Auth-Application-Id' = ?APP_ID},
+    {rejected, ?TIMEOUT} = call(?CLIENT2, Req).
+send_discard_2(_Config) ->
+    Req = #diameter_base_STR{'Destination-Realm' = realm(?SERVER4),
+                             'Termination-Cause' = ?MOVED,
+                             'Auth-Application-Id' = ?APP_ID},
+    {discarded, ?MOVED} = call(?CLIENT2, Req).
 
 %% ===========================================================================
 
 realm(Host) ->
     tl(lists:dropwhile(fun(C) -> C /= $. end, Host)).
 
-call(Req, Opts) ->
-    diameter:call(?CLIENT, ?APP_ALIAS, Req, Opts).
-
-set([H|T], Vs) ->
-    [H | Vs ++ T].
+call(Svc, Req) ->
+    diameter:call(Svc, ?APP_ALIAS, Req, [{filter, realm}]).
 
 %% ===========================================================================
 %% diameter callbacks
@@ -184,7 +205,8 @@ set([H|T], Vs) ->
 %% pick_peer/4
 
 %% Choose a server other than SERVER3 or SERVER5 if possible.
-pick_peer(Peers, _, ?CLIENT, _State) ->
+pick_peer(Peers, _, Svc, _State)
+  when ?IS_CLIENT(Svc) ->
     case lists:partition(fun({_, #diameter_caps{origin_host = {_, OH}}}) ->
                                  OH /= ?SERVER3 andalso OH /= ?SERVER5
                          end,
@@ -198,20 +220,41 @@ pick_peer(Peers, _, ?CLIENT, _State) ->
 
 %% prepare_request/3
 
-prepare_request(Pkt, ?CLIENT, {_Ref, Caps}) ->
+prepare_request(Pkt, Svc, {_Ref, Caps})
+  when ?IS_CLIENT(Svc) ->
     {send, prepare(Pkt, Caps)}.
 
 prepare(#diameter_packet{msg = Req}, Caps) ->
     #diameter_caps{origin_host  = {OH, _},
                    origin_realm = {OR, _}}
         = Caps,
-    set(Req, [{'Session-Id', diameter:session_id(OH)},
-              {'Origin-Host',  OH},
-              {'Origin-Realm', OR}]).
+    Req#diameter_base_STR{'Origin-Host' = OH,
+                          'Origin-Realm' = OR,
+                          'Session-Id' = diameter:session_id(OH)}.
+
+%% prepare_retransmit/3
+
+prepare_retransmit(#diameter_packet{header = H} = P, Svc, {_,_})
+  when ?IS_CLIENT(Svc) ->
+    #diameter_header{is_retransmitted = true} = H,  %% assert
+    prepare(P).
+
+prepare(#diameter_packet{msg = M} = P) ->
+    case M#diameter_base_STR.'Termination-Cause' of
+        ?LOGOUT  -> {send, P};
+        ?MOVED   -> discard;
+        ?TIMEOUT -> {discard, rejected}
+    end.
+
+%% handle_error/4
+
+handle_error(Reason, Req, _, _) ->
+    {Reason, Req#diameter_base_STR.'Termination-Cause'}.
 
 %% handle_answer/4
 
-handle_answer(Pkt, _Req, ?CLIENT, _Peer) ->
+handle_answer(Pkt, _Req, Svc, _Peer)
+  when ?IS_CLIENT(Svc) ->
     #diameter_packet{msg = Rec, errors = []} = Pkt,
     Rec.
 
@@ -219,8 +262,8 @@ handle_answer(Pkt, _Req, ?CLIENT, _Peer) ->
 
 %% Only SERVER3 actually answers.
 handle_request(Pkt, ?SERVER3, {_, Caps}) ->
-    #diameter_packet{msg = #diameter_base_STR{'Session-Id' = SId,
-                                              'Origin-Host' = ?CLIENT}}
+    #diameter_packet{header = #diameter_header{is_retransmitted = true},
+                     msg = #diameter_base_STR{'Session-Id' = SId}}
         = Pkt,
     #diameter_caps{origin_host  = {OH, _},
                    origin_realm = {OR, _}}