1 files changed, 33 insertions, 22 deletions
diff --git a/lib/diameter/examples/code/relay_cb.erl b/lib/diameter/examples/code/relay_cb.erl
index 6df1738143..0d56a14401 100644
--- a/lib/diameter/examples/code/relay_cb.erl
+++ b/lib/diameter/examples/code/relay_cb.erl
@@ -1,7 +1,7 @@
 %%
 %% %CopyrightBegin%
 %%
-%% Copyright Ericsson AB 2010-2016. All Rights Reserved.
+%% Copyright Ericsson AB 2010-2020. All Rights Reserved.
 %%
 %% Licensed under the Apache License, Version 2.0 (the "License");
 %% you may not use this file except in compliance with the License.
@@ -21,48 +21,59 @@
 -module(relay_cb).
 
 -include_lib("diameter/include/diameter.hrl").
--include_lib("diameter/include/diameter_gen_base_rfc3588.hrl").
 
 %% diameter callbacks
 -export([peer_up/3,
          peer_down/3,
-         pick_peer/5,
-         prepare_request/4,
-         prepare_retransmit/4,
-         handle_answer/5,
-         handle_error/5,
+         pick_peer/4,
+         prepare_request/3,
+         prepare_retransmit/3,
+         handle_answer/4,
+         handle_error/4,
          handle_request/3]).
 
+%% peer_up/3
+
 peer_up(_SvcName, _Peer, State) ->
     State.
 
+%% peer_down/3
+
 peer_down(_SvcName, _Peer, State) ->
     State.
 
-%% Returning 'relay' from handle_request causes diameter to resend the
-%% incoming request, which leads to pick_peer and prepare_request
-%% callbacks as if sending explicitly. The 'extra' argument is
-%% appended to the argument list for callbacks following from
-%% resending of the request.
+%% handle_request/3
+
+%% Assume the destination is directly connected; filter
+%% correspondingly; don't relay to the sender.
+handle_request(_Pkt, _SvcName, {_, Caps}) ->
+    #diameter_caps{origin_host = {_, OH}}
+        = Caps,
+    {relay, [{timeout, 2000},
+             {filter, {all, [host, realm, {neg, {host, OH}}]}}]}.
 
-handle_request(_Pkt, _SvcName, _Peer) ->
-    {relay, [{timeout, 1000}, {extra, [relayed]}]}.
+%% pick_peer/4
 
-%% diameter will filter the sender in the Peers list.
-pick_peer([Peer | _], _, _SvcName, _State, relayed) ->
+pick_peer([Peer | _], _, _SvcName, _State) ->
     {ok, Peer}.
 
-prepare_request(Pkt, _SvcName, _Peer, relayed) ->
+%% prepare_request/3
+
+prepare_request(Pkt, _SvcName, _Peer) ->
     {send, Pkt}.
 
-prepare_retransmit(Pkt, _SvcName, _Peer, relayed) ->
+%% prepare_request/3
+
+prepare_retransmit(Pkt, _SvcName, _Peer) ->
     {send, Pkt}.
 
-%% diameter expects handle_answer to return the diameter_packet record
-%% containing the answer when called for a relayed request.
+%% handle_answer/4
 
-handle_answer(Pkt, _Request, _SvcName, _Peer, relayed) ->
+%% Relay an answer by returning the first argument.
+handle_answer(Pkt, _Request, _SvcName, _Peer) ->
     Pkt.
 
-handle_error(Reason, _Request, _SvcName, _Peer, relayed) ->
+%% handle_error/4
+
+handle_error(Reason, _Request, _SvcName, _Peer) ->
     {error, Reason}.