From 83e11b6e39b833bd59dbc99e7b0cb6d6f41d6727 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=91=A8=E5=AD=90=E5=8D=9A?= <349832309@qq.com>
Date: Thu, 30 Aug 2018 22:45:08 +0800
Subject: [PATCH] Add emqx_banned test suite, and fix bugs in emqx_banned

---
 Makefile                   |  2 +-
 include/emqx.hrl           | 11 ++++++++++
 src/emqx_banned.erl        | 11 +++-------
 test/emqx_banned_SUITE.erl | 41 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 56 insertions(+), 9 deletions(-)
 create mode 100644 test/emqx_banned_SUITE.erl

diff --git a/Makefile b/Makefile
index aba0d4ac4..41ed22dc0 100644
--- a/Makefile
+++ b/Makefile
@@ -35,7 +35,7 @@ EUNIT_OPTS = verbose
 # CT_SUITES = emqx_mqueue
 ## emqx_trie emqx_router emqx_frame emqx_mqtt_compat
 
-CT_SUITES = emqx emqx_connection emqx_session emqx_access emqx_base62 emqx_broker emqx_client emqx_cm emqx_frame emqx_guid emqx_inflight \
+CT_SUITES = emqx emqx_banned emqx_connection emqx_session emqx_access emqx_base62 emqx_broker emqx_client emqx_cm emqx_frame emqx_guid emqx_inflight \
 			emqx_json emqx_keepalive emqx_lib emqx_metrics emqx_misc emqx_mod emqx_mqtt_caps \
 			emqx_mqtt_compat emqx_mqtt_properties emqx_mqueue emqx_message emqx_net emqx_pqueue emqx_router emqx_sm \
 			emqx_stats emqx_tables emqx_time emqx_topic emqx_trie emqx_vm emqx_zone emqx_mountpoint
diff --git a/include/emqx.hrl b/include/emqx.hrl
index 10190a3f6..34e41b0a1 100644
--- a/include/emqx.hrl
+++ b/include/emqx.hrl
@@ -145,5 +145,16 @@
           descr     :: string()
         }).
 
+%%--------------------------------------------------------------------
+%% Banned
+%%--------------------------------------------------------------------
+
+-record(banned, {
+          key, 
+          reason, 
+          by, 
+          desc, 
+          until}).
+
 -endif.
 
diff --git a/src/emqx_banned.erl b/src/emqx_banned.erl
index 908c8b5d5..4f8d44f44 100644
--- a/src/emqx_banned.erl
+++ b/src/emqx_banned.erl
@@ -36,14 +36,8 @@
 -define(TAB, ?MODULE).
 -define(SERVER, ?MODULE).
 
--type(key() :: {client_id, emqx_types:client_id()} |
-               {username, emqx_types:username() |
-               {ipaddr, inet:ip_address()}}).
-
 -record(state, {expiry_timer}).
 
--record(banned, {key :: key(), reason, by, desc, until}).
-
 %%--------------------------------------------------------------------
 %% Mnesia bootstrap
 %%--------------------------------------------------------------------
@@ -84,7 +78,7 @@ del(Key) ->
 %%--------------------------------------------------------------------
 
 init([]) ->
-    emqx_timer:seed(),
+    emqx_time:seed(),
     {ok, ensure_expiry_timer(#state{})}.
 
 handle_call(Req, _From, State) ->
@@ -128,7 +122,8 @@ expire_banned_item(Key, Now) ->
         [#banned{until = undefined}] -> ok;
         [B = #banned{until = Until}] when Until < Now ->
            mnesia:delete_object(?TAB, B, sticky_write);
+        [_] -> ok;
         [] -> ok
     end,
-    expire_banned_item(mnesia:next(Key), Now).
+    expire_banned_item(mnesia:next(?TAB, Key), Now).
 
diff --git a/test/emqx_banned_SUITE.erl b/test/emqx_banned_SUITE.erl
new file mode 100644
index 000000000..9fae880d4
--- /dev/null
+++ b/test/emqx_banned_SUITE.erl
@@ -0,0 +1,41 @@
+%% Copyright (c) 2018 EMQ Technologies Co., Ltd. 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.
+%% You may obtain a copy of the License at
+%%
+%%     http://www.apache.org/licenses/LICENSE-2.0
+%%
+%% Unless required by applicable law or agreed to in writing, software
+%% distributed under the License is distributed on an "AS IS" BASIS,
+%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+%% See the License for the specific language governing permissions and
+%% limitations under the License.
+
+-module(emqx_banned_SUITE).
+
+-compile(export_all).
+-compile(nowarn_export_all).
+
+-include("emqx.hrl").
+
+-include("emqx_mqtt.hrl").
+
+-include_lib("eunit/include/eunit.hrl").
+
+all() -> [t_banned_all].
+
+t_banned_all(_) ->
+    emqx_ct_broker_helpers:run_setup_steps(),
+    emqx_banned:start_link(),
+    {MegaSecs, Secs, MicroSecs} = erlang:timestamp(),
+    ok = emqx_banned:add(#banned{key = {client_id, <<"TestClient">>}, 
+                                 reason = <<"test">>,
+                                 by = <<"banned suite">>,
+                                 desc = <<"test">>, 
+                                 until = {MegaSecs, Secs + 10, MicroSecs}}),
+    % here is not expire banned test because its check interval is greater than 5 mins, but its effect has been confirmed
+    timer:sleep(100),
+    ?assert(emqx_banned:check(#{client_id => <<"TestClient">>, username => undefined, peername => {undefined, undefined}})),
+    emqx_banned:del({client_id, <<"TestClient">>}),
+    ?assertNot(emqx_banned:check(#{client_id => <<"TestClient">>, username => undefined, peername => {undefined, undefined}})).
\ No newline at end of file