From 80621ec9bef0ddfb01d5f037efb0ab04f2413cc9 Mon Sep 17 00:00:00 2001
From: Feng Lee <feng@emqx.io>
Date: Tue, 10 Sep 2019 23:38:24 +0800
Subject: [PATCH] Fix the test cases for emqx_banned, emqx_flapping modules

- Remove the 'is_enabled/1' function of 'emqx_banned' module
- Remove the 'is_enabled/1' function of 'emqx_flapping' module
- Add 'enable_acl/1', 'enable_banned/1' and 'enable_flapping_detect/1'
---
 src/emqx_banned.erl   | 15 ++++-----------
 src/emqx_channel.erl  | 11 ++++++-----
 src/emqx_flapping.erl | 13 ++++---------
 src/emqx_zone.erl     | 22 ++++++++++++++++++----
 4 files changed, 32 insertions(+), 29 deletions(-)

diff --git a/src/emqx_banned.erl b/src/emqx_banned.erl
index 4c1d7e969..c24bb4294 100644
--- a/src/emqx_banned.erl
+++ b/src/emqx_banned.erl
@@ -38,8 +38,6 @@
         , info/1
         ]).
 
--export([is_enabled/1]).
-
 %% gen_server callbacks
 -export([ init/1
         , handle_call/3
@@ -72,18 +70,13 @@ start_link() ->
     gen_server:start_link({local, ?MODULE}, ?MODULE, [], []).
 
 -spec(check(emqx_types:client()) -> boolean()).
-check(#{zone      := Zone,
-        client_id := ClientId,
+check(#{client_id := ClientId,
         username  := Username,
         peername  := {IPAddr, _}
        }) ->
-    is_enabled(Zone) andalso
-        ets:member(?BANNED_TAB, {client_id, ClientId})
-            orelse ets:member(?BANNED_TAB, {username, Username})
-                orelse ets:member(?BANNED_TAB, {ipaddr, IPAddr}).
-
-is_enabled(Zone) ->
-    emqx_zone:get_env(Zone, enable_ban, false).
+    ets:member(?BANNED_TAB, {client_id, ClientId})
+        orelse ets:member(?BANNED_TAB, {username, Username})
+            orelse ets:member(?BANNED_TAB, {ipaddr, IPAddr}).
 
 -spec(add(emqx_types:banned()) -> ok).
 add(Banned) when is_record(Banned, banned) ->
diff --git a/src/emqx_channel.erl b/src/emqx_channel.erl
index 2c2b7c02f..05f22edfb 100644
--- a/src/emqx_channel.erl
+++ b/src/emqx_channel.erl
@@ -964,14 +964,15 @@ set_logger_meta(_ConnPkt, #channel{client = #{client_id := ClientId}}) ->
 %% Check banned/flapping
 %%--------------------------------------------------------------------
 
-check_banned(_ConnPkt, #channel{client = Client}) ->
-    case emqx_banned:check(Client) of
+check_banned(_ConnPkt, #channel{client = Client = #{zone := Zone}}) ->
+    case emqx_zone:enable_banned(Zone) andalso emqx_banned:check(Client) of
         true  -> {error, ?RC_BANNED};
         false -> ok
     end.
 
-check_flapping(_ConnPkt, #channel{client = Client}) ->
-    case emqx_flapping:check(Client) of
+check_flapping(_ConnPkt, #channel{client = Client = #{zone := Zone}}) ->
+    case emqx_zone:enable_flapping_detect(Zone)
+         andalso emqx_flapping:check(Client) of
         true -> {error, ?RC_CONNECTION_RATE_EXCEEDED};
         false -> ok
     end.
@@ -1186,7 +1187,7 @@ maybe_resume_session(#channel{session  = Session,
 
 %% @doc Is ACL enabled?
 is_acl_enabled(#{zone := Zone, is_superuser := IsSuperuser}) ->
-    (not IsSuperuser) andalso emqx_zone:get_env(Zone, enable_acl, true).
+    (not IsSuperuser) andalso emqx_zone:enable_acl(Zone).
 
 %% @doc Parse Topic Filters
 -compile({inline, [parse_topic_filters/1]}).
diff --git a/src/emqx_flapping.erl b/src/emqx_flapping.erl
index 23a2696de..6e5f98c14 100644
--- a/src/emqx_flapping.erl
+++ b/src/emqx_flapping.erl
@@ -70,8 +70,8 @@ stop() -> gen_server:stop(?MODULE).
 
 %% @doc Check flapping when a MQTT client connected.
 -spec(check(emqx_types:client()) -> boolean()).
-check(#{zone := Zone, client_id := ClientId}) ->
-    is_enabled(Zone) andalso check(ClientId, get_policy()).
+check(#{client_id := ClientId}) ->
+    check(ClientId, get_policy()).
 
 check(ClientId, #{banned_interval := Interval}) ->
     case ets:lookup(?FLAPPING_TAB, {banned, ClientId}) of
@@ -82,8 +82,7 @@ check(ClientId, #{banned_interval := Interval}) ->
 
 %% @doc Detect flapping when a MQTT client disconnected.
 -spec(detect(emqx_types:client()) -> boolean()).
-detect(Client = #{zone := Zone}) ->
-    is_enabled(Zone) andalso detect(Client, get_policy()).
+detect(Client) -> detect(Client, get_policy()).
 
 detect(#{client_id := ClientId, peername := Peername},
        Policy = #{threshold := Threshold}) ->
@@ -107,11 +106,7 @@ detect(#{client_id := ClientId, peername := Peername},
             false
     end.
 
--compile({inline, [is_enabled/1, get_policy/0, now_diff/1]}).
-
--spec(is_enabled(emqx_types:zone()) -> boolean()).
-is_enabled(Zone) ->
-    emqx_zone:get_env(Zone, enable_flapping_detect, false).
+-compile({inline, [get_policy/0, now_diff/1]}).
 
 get_policy() ->
     emqx:get_env(flapping_detect_policy, ?DEFAULT_DETECT_POLICY).
diff --git a/src/emqx_zone.erl b/src/emqx_zone.erl
index 7a93037b1..586fa7ad7 100644
--- a/src/emqx_zone.erl
+++ b/src/emqx_zone.erl
@@ -25,7 +25,12 @@
 -logger_header("[Zone]").
 
 %% APIs
--export([start_link/0]).
+-export([start_link/0, stop/0]).
+
+-export([ enable_acl/1
+        , enable_banned/1
+        , enable_flapping_detect/1
+        ]).
 
 -export([ get_env/2
         , get_env/3
@@ -34,9 +39,6 @@
         , force_reload/0
         ]).
 
-%% for test
--export([stop/0]).
-
 %% gen_server callbacks
 -export([ init/1
         , handle_call/3
@@ -65,6 +67,18 @@
 start_link() ->
     gen_server:start_link({local, ?SERVER}, ?MODULE, [], []).
 
+-spec(enable_acl(zone()) -> boolean()).
+enable_acl(Zone) ->
+    get_env(Zone, enable_acl, true).
+
+-spec(enable_banned(zone()) -> boolean()).
+enable_banned(Zone) ->
+    get_env(Zone, enable_banned, false).
+
+-spec(enable_flapping_detect(zone()) -> boolean()).
+enable_flapping_detect(Zone) ->
+    get_env(Zone, enable_flapping_detect, false).
+
 -spec(get_env(maybe(zone()), atom()) -> maybe(term())).
 get_env(undefined, Key) -> emqx:get_env(Key);
 get_env(Zone, Key) ->