Przeglądaj źródła

Update for banned API

Use `mnesia:foldl` to traverse mnesia rather than `mnesia:first` and
`mnesia:next`, as a badarg exception would occur if the record was
deleted while travering the whole table.
terry-xiaoyu 7 lat temu
rodzic
commit
dd8513ad35
3 zmienionych plików z 20 dodań i 26 usunięć
  1. 9 6
      include/emqx.hrl
  2. 6 15
      src/emqx_banned.erl
  3. 5 5
      test/emqx_banned_SUITE.erl

+ 9 - 6
include/emqx.hrl

@@ -148,13 +148,16 @@
 %%--------------------------------------------------------------------
 %% Banned
 %%--------------------------------------------------------------------
+-type(banned_who() ::  {client_id,  binary()}
+                     | {username,   binary()}
+                     | {ip_address, inet:ip_address()}).
 
 -record(banned, {
-          key, 
-          reason, 
-          by, 
-          desc, 
-          until}).
+          who    :: banned_who(),
+          reason :: binary(),
+          by     :: binary(),
+          desc   :: binary(),
+          until  :: integer()
+        }).
 
 -endif.
-

+ 6 - 15
src/emqx_banned.erl

@@ -85,7 +85,7 @@ handle_cast(Msg, State) ->
     {noreply, State}.
 
 handle_info({timeout, TRef, expire}, State = #{expiry_timer := TRef}) ->
-    mnesia:async_dirty(fun expire_banned_items/1, [erlang:timestamp()]),
+    mnesia:async_dirty(fun expire_banned_items/1, [erlang:system_time(second)]),
     {noreply, ensure_expiry_timer(State), hibernate};
 
 handle_info(Info, State) ->
@@ -106,17 +106,8 @@ ensure_expiry_timer(State) ->
     State#{expiry_timer := emqx_misc:start_timer(timer:minutes(5), expire)}.
 
 expire_banned_items(Now) ->
-    expire_banned_item(mnesia:first(?TAB), Now).
-
-expire_banned_item('$end_of_table', _Now) ->
-    ok;
-expire_banned_item(Key, Now) ->
-    case mnesia:read(?TAB, Key) of
-        [#banned{until = undefined}] ->
-            ok;
-        [B = #banned{until = Until}] when Until < Now ->
-           mnesia:delete_object(?TAB, B, sticky_write);
-        _ -> ok
-    end,
-    expire_banned_item(mnesia:next(?TAB, Key), Now).
-
+    mnesia:foldl(fun
+            (B = #banned{until = Until}, _Acc) when Until < Now ->
+                mnesia:delete_object(?TAB, B, sticky_write);
+            (_, _Acc) -> ok
+        end, ok, ?TAB).

+ 5 - 5
test/emqx_banned_SUITE.erl

@@ -28,14 +28,14 @@ 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">>}, 
+    TimeNow = erlang:system_time(second),
+    ok = emqx_banned:add(#banned{who = {client_id, <<"TestClient">>},
                                  reason = <<"test">>,
                                  by = <<"banned suite">>,
-                                 desc = <<"test">>, 
-                                 until = {MegaSecs, Secs + 10, MicroSecs}}),
+                                 desc = <<"test">>,
+                                 until = TimeNow + 10}),
     % 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}})).
+    ?assertNot(emqx_banned:check(#{client_id => <<"TestClient">>, username => undefined, peername => {undefined, undefined}})).