Add incoming message interceptors

For plugins to provide interceptors they need to provide a cuttlefish
`.schema` file specifying the parameters for the provided interceptor.
Developers can take as example the existing mapping in `rabbit.schema`.
The order of the interceptors depends on the order of appearance of any
`message_interceptors.incoming.<interceptor>` entry in the
`rabbitmq.conf` file.

Co-authored-by: Péter Gömöri <[email protected]>

Author: LoisSotoLopez

deps/rabbit/Makefile

@@ -273,7 +273,7 @@ PARALLEL_CT_SET_3_B = cluster_upgrade list_consumers_sanity_check list_queues_on
 PARALLEL_CT_SET_3_C = cli_forget_cluster_node feature_flags_v2 mc_unit message_containers_deaths_v2 message_size_limit metadata_store_migration
 PARALLEL_CT_SET_3_D = metadata_store_phase1 metrics mirrored_supervisor peer_discovery_classic_config proxy_protocol runtime_parameters unit_stats_and_metrics unit_supervisor2 unit_vm_memory_monitor
 
-PARALLEL_CT_SET_4_A = clustering_events rabbit_local_random_exchange rabbit_message_interceptor rabbitmq_4_0_deprecations unit_pg_local unit_plugin_directories unit_plugin_versioning unit_policy_validators unit_priority_queue
+PARALLEL_CT_SET_4_A = clustering_events rabbit_local_random_exchange rabbit_incoming_message_interceptor rabbitmq_4_0_deprecations unit_pg_local unit_plugin_directories unit_plugin_versioning unit_policy_validators unit_priority_queue
 PARALLEL_CT_SET_4_B = per_user_connection_tracking per_vhost_connection_limit rabbit_fifo_dlx_integration rabbit_fifo_int
 PARALLEL_CT_SET_4_C = msg_size_metrics unit_msg_size_metrics per_vhost_msg_store per_vhost_queue_limit priority_queue upgrade_preparation vhost
 PARALLEL_CT_SET_4_D = per_user_connection_channel_tracking product_info publisher_confirms_parallel queue_type rabbitmq_queues_cli_integration rabbitmqctl_integration rabbitmqctl_shutdown routing rabbit_amqqueue

deps/rabbit/ct.test.spec

@@ -115,7 +115,7 @@
 ,   rabbit_fifo_prop_SUITE
 ,   rabbit_fifo_v0_SUITE
 ,   rabbit_local_random_exchange_SUITE
-,   rabbit_message_interceptor_SUITE
+,   rabbit_incoming_message_interceptor_SUITE
 ,   rabbit_stream_coordinator_SUITE
 ,   rabbit_stream_sac_coordinator_SUITE
 ,   rabbitmq_4_0_deprecations_SUITE

deps/rabbit/priv/schema/rabbit.schema

@@ -2658,22 +2658,56 @@ end}.
 {mapping, "message_interceptors.incoming.$interceptor.overwrite", "rabbit.incoming_message_interceptors", [
     {datatype, {enum, [true, false]}}]}.
 
+% Pseudo-key to include the interceptor in the list of interceptors. 
+%  - If any other configuration is provided for the interceptor this
+%  configuration is not required. 
+%  - If no other configuration is provided, this one is required so that the
+%  interceptor gets invoked. 
+{mapping, "message_interceptors.incoming.$interceptor.enabled", "rabbit.incoming_message_interceptors", [
+    {datatype, {enum, [true]}}]}.
+
+{mapping, "message_interceptors.incoming.set_header_timestamp.overwrite", "rabbit.incoming_message_interceptors", [
+    {datatype, {enum, [true, false]}}]}.
+
+{mapping, "message_interceptors.incoming.set_header_routing_node.overwrite", "rabbit.incoming_message_interceptors", [
+    {datatype, {enum, [true, false]}}]}.
+
 {translation, "rabbit.incoming_message_interceptors",
     fun(Conf) ->
-         case cuttlefish_variable:filter_by_prefix("message_interceptors", Conf) of
+         case cuttlefish_variable:filter_by_prefix("message_interceptors.incoming", Conf) of
              [] ->
                  cuttlefish:unset();
              L ->
-                 [begin
-                      Interceptor = list_to_atom(Interceptor0),
-                      case lists:member(Interceptor, [set_header_timestamp,
-                                                      set_header_routing_node]) of
-                          true ->
-                              {Interceptor, Overwrite};
-                          false ->
-                              cuttlefish:invalid(io_lib:format("~p is invalid", [Interceptor]))
-                      end
-                  end || {["message_interceptors", "incoming", Interceptor0, "overwrite"], Overwrite} <- L]
+                InterceptorsConfig = [ 
+                    {Module0, Config, Value}
+                    || {["message_interceptors", "incoming", Module0, Config], Value} <- L
+                ],
+                {Result, Order0} = lists:foldl(
+                    fun({Interceptor0, Key0, Value}, {Acc, Order}) -> 
+                        Interceptor = list_to_atom(Interceptor0),
+                        Key = list_to_atom(Key0),
+                        MapPutFun = fun(Key, Value) -> fun(Old) -> maps:put(Key, Value, Old) end end,
+                        Module = case Interceptor  of
+                            set_header_timestamp ->
+                                rabbit_header_timestamp_interceptor;
+                            set_header_routing_node ->
+                                rabbit_header_routing_node_interceptor;
+                            _ ->
+                                Interceptor
+                        end,
+                        NewAcc = 
+                            maps:update_with(
+                                Module, 
+                                MapPutFun(Key, Value), 
+                                #{Key => Value},
+                                Acc),
+                        {NewAcc, [Module| Order]}
+                    end,
+                    {#{}, []},
+                    InterceptorsConfig
+                ),
+                Order = lists:uniq(Order0),
+                [{O, maps:without([enabled], maps:get(O, Result))} || O <- Order]
          end
     end
 }.

deps/rabbit/src/rabbit_amqp_session.erl

@@ -10,6 +10,7 @@
 -compile({inline, [maps_update_with/4]}).
 
 -behaviour(gen_server).
+-behaviour(rabbit_protocol_accessor).
 
 -include_lib("kernel/include/logger.hrl").
 -include_lib("rabbit_common/include/rabbit.hrl").
@@ -105,6 +106,9 @@
          handle_info/2,
          format_status/1]).
 
+% `rabbit_protocol_accessor` behaviour callbacks
+-export([get_property/2]).
+
 -import(rabbit_amqp_util,
         [protocol_error/3]).
 -import(serial_number,
@@ -2436,7 +2440,7 @@ incoming_link_transfer(
     Mc0 = mc:init(mc_amqp, PayloadBin, #{}),
     case lookup_target(LinkExchange, LinkRKey, Mc0, Vhost, User, PermCache0) of
         {ok, X, RoutingKeys, Mc1, PermCache} ->
-            Mc2 = rabbit_message_interceptor:intercept(Mc1),
+            Mc2 = rabbit_incoming_message_interceptor:intercept(Mc1, ?MODULE, State0),
             check_user_id(Mc2, User),
             TopicPermCache = check_write_permitted_on_topics(
                                X, User, RoutingKeys, TopicPermCache0),
@@ -3874,6 +3878,15 @@ format_status(
               topic_permission_cache => TopicPermissionCache},
     maps:update(state, State, Status).
 
+get_property(user, #state{cfg = #cfg{user = User}}) ->
+    User; 
+get_property(vhost, #state{cfg = #cfg{vhost = VHost}}) ->
+    VHost; 
+get_property(connection_name, #state{cfg = #cfg{conn_name = ConnectionName}}) ->
+    ConnectionName;
+get_property(_, _) ->
+    undefined.
+
 -spec info(pid()) ->
     {ok, rabbit_types:infos()} | {error, term()}.
 info(Pid) ->

deps/rabbit/src/rabbit_channel.erl

@@ -7,6 +7,8 @@
 
 -module(rabbit_channel).
 
+-behaviour(rabbit_protocol_accessor).
+
 %% rabbit_channel processes represent an AMQP 0-9-1 channels.
 %%
 %% Connections parse protocol frames coming from clients and
@@ -60,6 +62,9 @@
          prioritise_call/4, prioritise_cast/3, prioritise_info/3,
          format_message_queue/2]).
 
+% `rabbit_protocol_accessor` behaviour callbacks
+-export([get_property/2]).
+
 %% Internal
 -export([list_local/0, emit_info_local/3, deliver_reply_local/3]).
 -export([get_vhost/1, get_user/1]).
@@ -813,6 +818,16 @@ get_consumer_timeout() ->
         _ ->
             undefined
     end.
+
+get_property(user, #ch{cfg = #conf{user = User}}) ->
+    User; 
+get_property(vhost, #ch{cfg = #conf{virtual_host = VHost}}) ->
+    VHost; 
+get_property(connection_name, #ch{cfg = #conf{conn_name = ConnectionName}}) ->
+    ConnectionName;
+get_property(_, _) ->
+    undefined.
+
 %%---------------------------------------------------------------------------
 
 reply(Reply, NewState) -> {reply, Reply, next_state(NewState), hibernate}.
@@ -1206,7 +1221,7 @@ handle_method(#'basic.publish'{exchange    = ExchangeNameBin,
             rabbit_misc:precondition_failed("invalid message: ~tp", [Reason]);
         {ok, Message0} ->
             check_write_permitted_on_topics(Exchange, User, Message0, AuthzContext),
-            Message = rabbit_message_interceptor:intercept(Message0),
+            Message = rabbit_incoming_message_interceptor:intercept(Message0, ?MODULE, State),
             check_user_id_header(Message, User),
             QNames = rabbit_exchange:route(Exchange, Message, #{return_binding_keys => true}),
             [deliver_reply(RK, Message) || {virtual_reply_queue, RK} <- QNames],

deps/rabbit/src/rabbit_header_routing_node_interceptor.erl

@@ -0,0 +1,17 @@
+-module(rabbit_header_routing_node_interceptor).
+-behaviour(rabbit_incoming_message_interceptor).
+
+-define(HEADER_ROUTING_NODE, <<"x-routed-by">>).
+
+-export([
+    intercept/4
+]).
+
+intercept(Msg, _ProtoMod, _ProtoState, Config) ->
+    Node = atom_to_binary(node()),
+    Overwrite = maps:get(overwrite, Config, false),
+    rabbit_incoming_message_interceptor:set_msg_annotation(
+        Msg,
+        ?HEADER_ROUTING_NODE,
+        Node,
+        Overwrite).

deps/rabbit/src/rabbit_header_timestamp_interceptor.erl

@@ -0,0 +1,28 @@
+-module(rabbit_header_timestamp_interceptor).
+-behaviour(rabbit_incoming_message_interceptor).
+
+-include("mc.hrl").
+
+-define(HEADER_TIMESTAMP, <<"timestamp_in_ms">>).
+
+-export([
+    intercept/4
+]).
+
+intercept(Msg0, _ProtoMod, _ProtoState, Config) ->
+    Ts = mc:get_annotation(?ANN_RECEIVED_AT_TIMESTAMP, Msg0),
+    Overwrite = maps:get(overwrite, Config, false),
+    Msg = rabbit_incoming_message_interceptor:set_msg_annotation(
+        Msg0,
+        ?HEADER_TIMESTAMP,
+        Ts,
+        Overwrite),
+    set_msg_timestamp(Msg, Ts, Overwrite).
+
+set_msg_timestamp(Msg, Timestamp, Overwrite) ->
+    case {mc:timestamp(Msg), Overwrite} of
+        {Ts, false} when is_integer(Ts) ->
+            Msg;
+        _ ->
+            mc:set_annotation(?ANN_TIMESTAMP, Timestamp, Msg)
+    end.

deps/rabbit/src/rabbit_incoming_message_interceptor.erl

@@ -0,0 +1,38 @@
+-module(rabbit_incoming_message_interceptor).
+
+-export([
+    intercept/3,
+    set_msg_annotation/4
+]).
+
+-callback intercept(
+    Msg :: mc:state(),
+    ProtoMod :: module(),
+    ProtoState :: term(),
+    Config :: #{atom() := term()}
+) -> mc:state().
+
+-spec intercept(Msg, ProtoMod, ProtoState) -> Resp when
+    Msg :: mc:state(),
+    ProtoMod :: module(),
+    ProtoState :: term(),
+    Resp :: mc:state().
+intercept(Msg, ProtoMod, ProtoState) ->
+    Interceptors = persistent_term:get(incoming_message_interceptors, []),
+    lists:foldl(fun({Module, Config}, Msg0) ->
+        try
+            Module:intercept(Msg0, ProtoMod, ProtoState, Config)
+        catch
+            error:undef ->
+                Msg0
+        end
+    end, Msg , Interceptors).
+
+-spec set_msg_annotation(mc:state(), mc:ann_key(), mc:ann_value(), boolean()) -> mc:state().
+set_msg_annotation(Msg, Key, Value, Overwrite) ->
+    case {mc:x_header(Key, Msg), Overwrite} of
+        {Val, false} when Val =/= undefined ->
+            Msg;
+        _ ->
+            mc:set_annotation(Key, Value, Msg)
+    end.

deps/rabbit/src/rabbit_message_interceptor.erl

@@ -0,0 +1,27 @@
+-module(rabbit_protocol_accessor).
+
+-callback get_property(Property, State) -> Resp when
+    Property :: atom(),
+    State :: term(),
+    Resp :: term().
+
+-export([
+    get_proto_state_properties/3
+]).
+
+-spec get_proto_state_properties(
+    ProtoAccessor :: module(),
+    ProtoState :: term(),
+    Properties :: [Property]
+) -> Resp :: [{Property, term()}] when
+    Property :: atom().
+get_proto_state_properties(ProtoAccessor, ProtoState, Properties) ->
+    lists:foldl(
+        fun(Property, Results) ->
+            case ProtoAccessor:get_property(Property, ProtoState) of
+                undefined -> Results;
+                Value -> [{Property, Value} | Results]
+            end
+        end,
+        [],
+        Properties).