Don't count messages on arbitrary keys

That is, don't use a key constructed from an incoming Diameter header
unless the message is known to the dictionary in question.  Otherwise
there are 2^32 application ids, 2^24 command codes, and 2 R-bits for an
ill-willed peer to choose from, each resulting in new keys in the
counter table (diameter_stats).

The usual {ApplicationId, CommandCode, Rbit} in a key is replaced by the
atom 'unknown' if the message in question is unknown to the decoding
dictionary.

Counters for messages sent and received by a relay are (still) not
implemented.

3 files changed, 102 insertions, 53 deletions

diff --git a/lib/diameter/src/base/diameter_peer_fsm.erl b/lib/diameter/src/base/diameter_peer_fsm.erl
index e9b6a263c9..58166349f0 100644
--- a/lib/diameter/src/base/diameter_peer_fsm.erl
+++ b/lib/diameter/src/base/diameter_peer_fsm.erl
@@ -293,7 +293,7 @@ handle_info(T, #state{} = State) ->
             {stop, {shutdown, T}, State}
     catch
         exit: {diameter_codec, encode, T} = Reason ->
-            incr_error(send, T),
+            incr_error(send, T, State#state.dictionary),
             ?LOG(stop, Reason),
             {stop, {shutdown, Reason}, State};
         {?MODULE, Tag, Reason}  ->
@@ -523,7 +523,6 @@ recv(#diameter_packet{header = #diameter_header{} = Hdr}
      = S) ->
     Name = diameter_codec:msg_name(Dict0, Hdr),
     Pid ! {recv, self(), Name, Pkt},
-    diameter_stats:incr({msg_id(Name, Hdr), recv}), %% count received
     rcv(Name, Pkt, S);
 
 recv(#diameter_packet{header = undefined,
@@ -564,20 +563,16 @@ invalid(E, Reason, T) ->
     ?LOG(Reason, T),
     ok.
 
-msg_id({_,_,_} = T, _) ->
-    T;
-msg_id(_, Hdr) ->
-    {_,_,_} = diameter_codec:msg_id(Hdr).
-
 %% rcv/3
 
 %% Incoming CEA.
-rcv('CEA',
+rcv('CEA' = N,
     #diameter_packet{header = #diameter_header{end_to_end_id = Eid,
                                                hop_by_hop_id = Hid}}
     = Pkt,
     #state{state = {'Wait-CEA', Hid, Eid}}
     = S) ->
+    ?LOG(recv, N),
     handle_CEA(Pkt, S);
 
 %% Incoming CER
@@ -598,29 +593,46 @@ rcv('DPR' = N, Pkt, S) ->
 %% DPA in response to DPR and with the expected identifiers.
 rcv('DPA' = N,
     #diameter_packet{header = #diameter_header{end_to_end_id = Eid,
-                                               hop_by_hop_id = Hid}}
+                                               hop_by_hop_id = Hid}
+                            = H}
     = Pkt,
     #state{dictionary = Dict0,
            transport = TPid,
            dpr = {Hid, Eid}}) ->
-
+    ?LOG(recv, N),
+    incr(recv, H, Dict0),
     incr_rc(recv, diameter_codec:decode(Dict0, Pkt), Dict0),
     diameter_peer:close(TPid),
     {stop, N};
 
 %% Ignore anything else, an unsolicited DPA in particular.
+rcv(N, #diameter_packet{header = H}, _)
+  when N == 'CER';
+       N == 'CEA';
+       N == 'DPR';
+       N == 'DPA' ->
+    ?LOG(ignored, N),
+    %% Note that these aren't counted in the normal recv counter.
+    diameter_stats:incr({diameter_codec:msg_id(H), recv, ignored}),
+    ok;
+
 rcv(_, _, _) ->
     ok.
 
+%% incr/3
+
+incr(Dir, Hdr, Dict0) ->
+    diameter_traffic:incr(Dir, Hdr, self(), Dict0).
+
 %% incr_rc/3
 
 incr_rc(Dir, Pkt, Dict0) ->
     diameter_traffic:incr_rc(Dir, Pkt, self(), Dict0).
 
-%% incr_error/2
+%% incr_error/3
 
-incr_error(Dir, Pkt) ->
-    diameter_traffic:incr_error(Dir, Pkt, self()).
+incr_error(Dir, Pkt, Dict0) ->
+    diameter_traffic:incr_error(Dir, Pkt, self(), Dict0).
 
 %% send/2
 
@@ -628,19 +640,23 @@ incr_error(Dir, Pkt) ->
 %% sending. In particular, the watchdog will send DWR as a binary
 %% while messages coming from clients will be in a #diameter_packet.
 send(Pid, Msg) ->
-    diameter_stats:incr({diameter_codec:msg_id(Msg), send}),
     diameter_peer:send(Pid, Msg).
 
 %% handle_request/3
+%%
+%% Incoming CER or DPR.
 
-handle_request(Type, #diameter_packet{} = Pkt, #state{dictionary = D} = S) ->
-    ?LOG(recv, Type),
-    send_answer(Type, diameter_codec:decode(D, Pkt), S).
+handle_request(Name,
+               #diameter_packet{header = H} = Pkt,
+               #state{dictionary = Dict0} = S) ->
+    ?LOG(recv, Name),
+    incr(recv, H, Dict0),
+    send_answer(Name, diameter_codec:decode(Dict0, Pkt), S).
 
 %% send_answer/3
 
 send_answer(Type, ReqPkt, #state{transport = TPid, dictionary = Dict} = S) ->
-    incr_error(recv, ReqPkt),
+    incr_error(recv, ReqPkt, Dict),
 
     #diameter_packet{header = H,
                      transport_data = TD}
@@ -660,6 +676,7 @@ send_answer(Type, ReqPkt, #state{transport = TPid, dictionary = Dict} = S) ->
 
     AnsPkt = diameter_codec:encode(Dict, Pkt),
 
+    incr(send, AnsPkt, Dict),
     incr_rc(send, AnsPkt, Dict),
     send(TPid, AnsPkt),
     ?LOG(send, ans(Type)),
@@ -862,13 +879,12 @@ recv_CER(CER, #state{service = Svc, dictionary = Dict}) ->
 
 %% handle_CEA/2
 
-handle_CEA(#diameter_packet{bin = Bin}
+handle_CEA(#diameter_packet{header = H}
            = Pkt,
            #state{dictionary = Dict0,
                   service = #diameter_service{capabilities = LCaps}}
-           = S)
-  when is_binary(Bin) ->
-    ?LOG(recv, 'CEA'),
+           = S) ->
+    incr(recv, H, Dict0),
 
     #diameter_packet{}
         = DPkt
@@ -909,7 +925,7 @@ handle_CEA(#diameter_packet{bin = Bin}
 %% capabilities exchange could send DIAMETER_LIMITED_SUCCESS = 2002,
 %% even if this isn't required by RFC 3588.
 
-result_code({{'Result-Code', N}, _}) ->
+result_code({'Result-Code', N}) ->
     N;
 result_code(_) ->
     undefined.
diff --git a/lib/diameter/src/base/diameter_traffic.erl b/lib/diameter/src/base/diameter_traffic.erl
index 7c680f96b5..5fac61f416 100644
--- a/lib/diameter/src/base/diameter_traffic.erl
+++ b/lib/diameter/src/base/diameter_traffic.erl
@@ -32,7 +32,8 @@
 -export([receive_message/4]).
 
 %% towards diameter_peer_fsm and diameter_watchdog
--export([incr_error/3,
+-export([incr/4,
+         incr_error/4,
          incr_rc/4]).
 
 %% towards diameter_service
@@ -115,38 +116,52 @@ peer_down(TPid) ->
     failover(TPid).
 
 %% ---------------------------------------------------------------------------
-%% incr_error/3
+%% incr/4
 %% ---------------------------------------------------------------------------
 
-%% A decoded message with errors.
-incr_error(Dir, #diameter_packet{header = H, errors = [_|_]}, TPid) ->
-    incr_error(Dir, H, TPid);
+incr(Dir, #diameter_packet{header = H}, TPid, Dict) ->
+    incr(Dir, H, TPid, Dict);
 
-%% An encoded message with errors and an identifiable header ...
-incr_error(Dir, {_, _, #diameter_header{} = H}, TPid) ->
-    incr_error(Dir, H, TPid);
+incr(Dir, #diameter_header{} = H, TPid, Dict) ->
+    incr(TPid, {msg_id(H, Dict), Dir}).
 
-%% ... or not.
-incr_error(Dir, {_,_}, TPid) ->
-    incr(TPid, {unknown, Dir, error});
+%% ---------------------------------------------------------------------------
+%% incr_error/4
+%% ---------------------------------------------------------------------------
 
-incr_error(Dir, #diameter_header{} = H, TPid) ->
-    incr_error(Dir, diameter_codec:msg_id(H), TPid);
+%% Decoded message without errors.
+incr_error(recv, #diameter_packet{errors = []}, _, _) ->
+    ok;
 
-incr_error(Dir, {_,_,_} = Id, TPid) ->
-    incr(TPid, {Id, Dir, error});
+incr_error(recv = D, #diameter_packet{header = H}, TPid, Dict) ->
+    incr_error(D, H, TPid, Dict);
 
-incr_error(_, _, _) ->
-    false.
+%% Encoded message with errors and an identifiable header ...
+incr_error(send = D, {_, _, #diameter_header{} = H}, TPid, Dict) ->
+    incr_error(D, H, TPid, Dict);
 
+%% ... or not.
+incr_error(send = D, {_,_}, TPid, _) ->
+    incr_error(D, unknown, TPid);
+
+incr_error(Dir, #diameter_header{} = H, TPid, Dict) ->
+    incr_error(Dir, msg_id(H, Dict), TPid);
+
+incr_error(Dir, Id, TPid, _) ->
+    incr_error(Dir, Id, TPid).
+
+incr_error(Dir, Id, TPid) ->
+    incr(TPid, {Id, Dir, error}).
+    
 %% ---------------------------------------------------------------------------
 %% incr_rc/4
 %% ---------------------------------------------------------------------------
 
--spec incr_rc(send|recv, #diameter_packet{}, TPid, Dict0)
+-spec incr_rc(send|recv, Pkt, TPid, Dict0)
    -> {Counter, non_neg_integer()}
     | Reason
- when TPid :: pid(),
+ when Pkt :: #diameter_packet{},
+      TPid :: pid(),
       Dict0 :: module(),
       Counter :: {'Result-Code', integer()}
                | {'Experimental-Result', integer(), integer()},
@@ -264,8 +279,9 @@ recv_R({#diameter_app{id = Id, dictionary = Dict} = App, Caps},
        Pkt0,
        Dict0,
        RecvData) ->
+    incr(recv, Pkt0, TPid, Dict),
     Pkt = errors(Id, diameter_codec:decode(Id, Dict, Pkt0)),
-    incr_error(recv, Pkt, TPid),
+    incr_error(recv, Pkt, TPid, Dict),
     {Caps, Pkt, App, recv_R(App, TPid, Dict0, Caps, RecvData, Pkt)};
 %% Note that the decode is different depending on whether or not Id is
 %% ?APP_ID_RELAY.
@@ -656,6 +672,7 @@ reply(Msg, Dict, TPid, Dict0, Fs, ReqPkt) ->
                  TPid,
                  reset(make_answer_packet(Msg, ReqPkt), Dict, Dict0),
                  Fs),
+    incr(send, Pkt, TPid, Dict),
     incr_rc(send, Pkt, Dict, TPid, Dict0),  %% count outgoing
     send(TPid, Pkt).
 
@@ -1040,10 +1057,10 @@ incr_rc(Dir, Pkt, Dict, TPid, Dict0) ->
                      errors = Es}
         = Pkt,
 
-    Id = diameter_codec:msg_id(Hdr),
+    Id = msg_id(Hdr, Dict),
 
     %% Count incoming decode errors.
-    recv /= Dir orelse [] == Es orelse incr_error(Dir, Id, TPid),
+    recv /= Dir orelse [] == Es orelse incr_error(Dir, Id, TPid, Dict),
 
     %% Exit on a missing result code.
     T = rc_counter(Dict, Msg),
@@ -1054,7 +1071,15 @@ incr_rc(Dir, Pkt, Dict, TPid, Dict0) ->
     is_result(RC, E, Dict0)
         orelse ?LOGX(invalid_error_bit, {Dict, Dir, Hdr, RC}),
 
-    incr(TPid, {Id, Dir, Ctr}).
+    incr(TPid, {Id, Dir, Ctr}),
+    Ctr.
+
+%% Only count on known keeps so as not to be vulnerable to attack:
+%% there are 2^32 (application ids) * 2^24 (command codes) * 2 (R-bits)
+%% = 2^57 Ids for an attacker to choose from.
+msg_id(Hdr, Dict) ->
+    {_ApplId, Code, R} = Id = diameter_codec:msg_id(Hdr),
+    choose('' == Dict:msg_name(Code, 0 == R), unknown, Id).
 
 %% No E-bit: can't be 3xxx.
 is_result(RC, false, _Dict0) ->
@@ -1072,8 +1097,8 @@ is_result(RC, true, _) ->
 
 %% incr/2
 
-incr(TPid, {_, _, T} = Counter) ->
-    {T, diameter_stats:incr(Counter, TPid, 1)}.
+incr(TPid, Counter) ->
+    diameter_stats:incr(Counter, TPid, 1).
 
 %% rc_counter/2
 
@@ -1423,6 +1448,8 @@ handle_answer(SvcName,
 %% want to examine the answer?
 
 handle_A(Pkt, SvcName, Dict, Dict0, App, #request{transport = TPid} = Req) ->
+    incr(recv, Pkt, TPid, Dict),
+
     try
         incr_rc(recv, Pkt, Dict, TPid, Dict0) %% count incoming
     of
@@ -1547,7 +1574,7 @@ encode(Dict, TPid, #diameter_packet{bin = undefined} = Pkt) ->
         diameter_codec:encode(Dict, Pkt)
     catch
         exit: {diameter_codec, encode, T} = Reason ->
-            incr_error(send, T, TPid),
+            incr_error(send, T, TPid, Dict),
             exit(Reason)
     end;
 
diff --git a/lib/diameter/src/base/diameter_watchdog.erl b/lib/diameter/src/base/diameter_watchdog.erl
index 5e27778432..4616c675ff 100644
--- a/lib/diameter/src/base/diameter_watchdog.erl
+++ b/lib/diameter/src/base/diameter_watchdog.erl
@@ -471,8 +471,7 @@ encode(dwr = M, Dict0, Mask) ->
                            hop_by_hop_id = Seq},
     Pkt = #diameter_packet{header = Hdr,
                            msg = Msg},
-    #diameter_packet{bin = Bin} = diameter_codec:encode(Dict0, Pkt),
-    Bin;
+    diameter_codec:encode(Dict0, Pkt);
 
 encode(dwa, Dict0, #diameter_packet{header = H, transport_data = TD}
                    = ReqPkt) ->
@@ -541,10 +540,14 @@ send_watchdog(#watchdog{pending = false,
                         dictionary = Dict0,
                         sequence = Mask}
               = S) ->
-    send(TPid, {send, encode(dwr, Dict0, Mask)}),
+    #diameter_packet{bin = Bin} = EPkt = encode(dwr, Dict0, Mask),
+    diameter_traffic:incr(send, EPkt, TPid, Dict0),
+    send(TPid, {send, Bin}),
     ?LOG(send, 'DWR'),
     S#watchdog{pending = true}.
 
+%% Dont' count encode errors since we don't expect any on DWR/DWA.
+
 %% recv/3
 
 recv(Name, Pkt, S) ->
@@ -563,8 +566,10 @@ rcv('DWR', Pkt, #watchdog{transport = TPid,
                           dictionary = Dict0}) ->
     ?LOG(recv, 'DWR'),
     DPkt = diameter_codec:decode(Dict0, Pkt),
-    diameter_traffic:incr_error(recv, DPkt, TPid),
+    diameter_traffic:incr(recv, DPkt, TPid, Dict0),
+    diameter_traffic:incr_error(recv, DPkt, TPid, Dict0),
     EPkt = encode(dwa, Dict0, Pkt),
+    diameter_traffic:incr(send, EPkt, TPid, Dict0),
     diameter_traffic:incr_rc(send, EPkt, TPid, Dict0),
 
     send(TPid, {send, EPkt}),
@@ -573,6 +578,7 @@ rcv('DWR', Pkt, #watchdog{transport = TPid,
 rcv('DWA', Pkt, #watchdog{transport = TPid,
                           dictionary = Dict0}) ->
     ?LOG(recv, 'DWA'),
+    diameter_traffic:incr(recv, Pkt, TPid, Dict0),
     diameter_traffic:incr_rc(recv,
                              diameter_codec:decode(Dict0, Pkt),
                              TPid,