Merge pull request #217 from velimir/cache-topic-access

cache topic access

(cherry picked from commit 1c9e9d7)

Author: michaelklishin

src/rabbit_mqtt_processor.erl

@@ -18,7 +18,7 @@
 
 -export([info/2, initial_state/2, initial_state/5,
          process_frame/2, amqp_pub/2, amqp_callback/2, send_will/1,
-         close_connection/1]).
+         close_connection/1, handle_pre_hibernate/0]).
 
 %% for testing purposes
 -export([get_vhost_username/1, get_vhost/3, get_vhost_from_user_mapping/2,
@@ -31,6 +31,7 @@
 -define(APP, rabbitmq_mqtt).
 -define(FRAME_TYPE(Frame, Type),
         Frame = #mqtt_frame{ fixed = #mqtt_frame_fixed{ type = Type }}).
+-define(MAX_TOPIC_PERMISSION_CACHE_SIZE, 12).
 
 initial_state(Socket, SSLLoginName) ->
     RealSocket = rabbit_net:unwrap_socket(Socket),
@@ -884,9 +885,19 @@ close_connection(PState = #proc_state{ connection = Connection,
     PState #proc_state{ channels   = {undefined, undefined},
                         connection = undefined }.
 
+<<<<<<< HEAD
 % NB: check_*: MQTT spec says we should ack normally, ie pretend there
 % was no auth error, but here we are closing the connection with an error. This
 % is what happens anyway if there is an authorization failure at the AMQP level.
+=======
+handle_pre_hibernate() ->
+    erase(topic_permission_cache),
+    ok.
+
+%% NB: check_*: MQTT spec says we should ack normally, ie pretend there
+%% was no auth error, but here we are closing the connection with an error. This
+%% is what happens anyway if there is an authorization failure at the AMQP 0-9-1 client level.
+>>>>>>> 1c9e9d7... Merge pull request #217 from velimir/cache-topic-access
 
 check_publish(TopicName, Fn, PState) ->
   case check_topic_access(TopicName, write, PState) of
@@ -910,28 +921,46 @@ check_topic_access(TopicName, Access,
                         exchange = Exchange,
                         client_id = ClientId,
                         mqtt2amqp_fun = Mqtt2AmqpFun }) ->
-  Resource = #resource{virtual_host = VHost,
-                       kind = topic,
-                       name = Exchange},
-  RoutingKey = Mqtt2AmqpFun(TopicName),
-  Context = #{routing_key  => RoutingKey,
-              variable_map => #{
-                  <<"username">>  => Username,
-                  <<"vhost">>     => VHost,
-                  <<"client_id">> => rabbit_data_coercion:to_binary(ClientId)
-              }
-  },
-
-  try rabbit_access_control:check_topic_access(User, Resource, Access, Context) of
-    R -> R
-  catch
-    _:{amqp_error, access_refused, Msg, _} ->
-      rabbit_log:error("operation resulted in an error (access_refused): ~p~n", [Msg]),
-        {error, access_refused};
-    _:Error ->
-      rabbit_log:error("~p~n", [Error]),
-        {error, access_refused}
-  end.
+    Cache =
+        case get(topic_permission_cache) of
+            undefined -> [];
+            Other     -> Other
+        end,
+
+    Key = {TopicName, Username, ClientId, VHost, Exchange, Access},
+    case lists:member(Key, Cache) of
+        true ->
+            ok;
+        false ->
+            Resource = #resource{virtual_host = VHost,
+                                 kind = topic,
+                                 name = Exchange},
+
+            RoutingKey = Mqtt2AmqpFun(TopicName),
+            Context = #{routing_key  => RoutingKey,
+                        variable_map => #{
+                                          <<"username">>  => Username,
+                                          <<"vhost">>     => VHost,
+                                          <<"client_id">> => rabbit_data_coercion:to_binary(ClientId)
+                                         }
+                       },
+
+            try rabbit_access_control:check_topic_access(User, Resource, Access, Context) of
+                ok ->
+                    CacheTail = lists:sublist(Cache, ?MAX_TOPIC_PERMISSION_CACHE_SIZE - 1),
+                    put(topic_permission_cache, [Key | CacheTail]),
+                    ok;
+                R ->
+                    R
+            catch
+                _:{amqp_error, access_refused, Msg, _} ->
+                    rabbit_log:error("operation resulted in an error (access_refused): ~p~n", [Msg]),
+                    {error, access_refused};
+                _:Error ->
+                    rabbit_log:error("~p~n", [Error]),
+                    {error, access_refused}
+            end
+    end.
 
 info(consumer_tags, #proc_state{consumer_tags = Val}) -> Val;
 info(unacked_pubs, #proc_state{unacked_pubs = Val}) -> Val;

src/rabbit_mqtt_reader.erl

@@ -24,7 +24,7 @@
 
 -export([start_link/2]).
 -export([init/1, handle_call/3, handle_cast/2, handle_info/2,
-         code_change/3, terminate/2]).
+         code_change/3, terminate/2, handle_pre_hibernate/1]).
 
 -export([conserve_resources/3, start_keepalive/2]).
 
@@ -189,6 +189,10 @@ terminate(Reason, State) ->
     maybe_emit_stats(State),
     do_terminate(Reason, State).
 
+handle_pre_hibernate(State) ->
+    rabbit_mqtt_processor:handle_pre_hibernate(),
+    {hibernate, State}.
+
 do_terminate({network_error, {ssl_upgrade_error, closed}, ConnStr}, _State) ->
     rabbit_log_connection:error("MQTT detected TLS upgrade error on ~s: connection closed~n",
        [ConnStr]);