Add incoming message interceptors

This commit enables users to provide custom message interceptor modules,
i.e. modules to process incoming and outgoing messages. The
`rabbit_message_interceptor` behaviour defines a `intercept/4` callback,
for those modules to implement.

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

Author: 2 people

deps/rabbit/priv/schema/rabbit.schema

@@ -2667,23 +2667,102 @@ 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.outgoing.$interceptor.enabled", "rabbit.outgoing_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.rabbit_message_interceptor_routing_node.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]}}]}.
+{mapping,
+    "message_interceptors.incoming.rabbit_message_interceptor_timestamp.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
-             [] ->
-                 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]
-         end
+        case cuttlefish_variable:filter_by_prefix("message_interceptors.incoming", Conf) of
+            [] ->
+                cuttlefish:unset();
+            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(Old) -> maps:put(Key, Value, Old) end,
+                        % This Interceptor -> Module alias exists for
+                        % compatibility reasons
+                        Module = case Interceptor  of
+                                     set_header_timestamp ->
+                                         rabbit_message_interceptor_timestamp;
+                                     set_header_routing_node ->
+                                         rabbit_message_interceptor_routing_node;
+                                     _ ->
+                                         Interceptor
+                                 end,
+                        NewAcc = maps:update_with(Module,
+                                                  MapPutFun,
+                                                  #{Key => Value},
+                                                  Acc),
+                        {NewAcc, [Module| Order]}
+                    end,
+                    {#{}, []},
+                    InterceptorsConfig
+                ),
+                Order = lists:uniq(Order0),
+                [{O, maps:without([enabled], maps:get(O, Result))} || O <- Order]
+        end
+    end
+}.
+
+{translation, "rabbit.outgoing_message_interceptors",
+    fun(Conf) ->
+        case cuttlefish_variable:filter_by_prefix("message_interceptors.outgoing", Conf) of
+            [] ->
+                cuttlefish:unset();
+            L ->
+                InterceptorsConfig = [
+                    {Module0, Config, Value}
+                    || {["message_interceptors", "outgoing", Module0, Config], Value} <- L
+                ],
+                {Result, Order0} = lists:foldl(
+                    fun({Interceptor0, Key0, Value}, {Acc, Order}) ->
+                        Module = list_to_atom(Interceptor0),
+                        Key = list_to_atom(Key0),
+                        MapPutFun = fun(Old) -> maps:put(Key, Value, Old) end,
+                        NewAcc = maps:update_with(Module,
+                                                  MapPutFun,
+                                                  #{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.erl

@@ -1656,7 +1656,8 @@ persist_static_configuration() ->
       [classic_queue_index_v2_segment_entry_count,
        classic_queue_store_v2_max_cache_size,
        classic_queue_store_v2_check_crc32,
-       incoming_message_interceptors
+       incoming_message_interceptors,
+       outgoing_message_interceptors
       ]),
 
     %% Disallow the following two cases:

deps/rabbit/src/rabbit_amqp_session.erl

@@ -283,7 +283,8 @@
           max_handle :: link_handle(),
           max_incoming_window :: pos_integer(),
           max_link_credit :: pos_integer(),
-          max_queue_credit :: pos_integer()
+          max_queue_credit :: pos_integer(),
+          msg_interceptor_ctx :: map()
          }).
 
 -record(state, {
@@ -474,7 +475,11 @@ init({ReaderPid, WriterPid, ChannelNum, MaxFrameSize, User, Vhost, ContainerId,
                            max_handle = EffectiveHandleMax,
                            max_incoming_window = MaxIncomingWindow,
                            max_link_credit = MaxLinkCredit,
-                           max_queue_credit = MaxQueueCredit
+                           max_queue_credit = MaxQueueCredit,
+                           msg_interceptor_ctx = #{protocol => ?PROTOCOL,
+                                                   username => User#user.username,
+                                                   vhost => Vhost,
+                                                   conn_name => ConnName}
                           }}}.
 
 terminate(_Reason, #state{incoming_links = IncomingLinks,
@@ -2411,7 +2416,8 @@ incoming_link_transfer(
                              trace_state = Trace,
                              conn_name = ConnName,
                              channel_num = ChannelNum,
-                             max_link_credit = MaxLinkCredit}}) ->
+                             max_link_credit = MaxLinkCredit,
+                             msg_interceptor_ctx = MsgInterceptorCtx}}) ->
 
     {PayloadBin, DeliveryId, Settled} =
     case MultiTransfer of
@@ -2436,7 +2442,9 @@ 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_message_interceptor:intercept(Mc1,
+                                                       MsgInterceptorCtx,
+                                                       incoming_message_interceptors),
             check_user_id(Mc2, User),
             TopicPermCache = check_write_permitted_on_topics(
                                X, User, RoutingKeys, TopicPermCache0),

deps/rabbit/src/rabbit_channel.erl

@@ -110,7 +110,8 @@
           authz_context,
           max_consumers,  % taken from rabbit.consumer_max_per_channel
           %% defines how ofter gc will be executed
-          writer_gc_threshold
+          writer_gc_threshold,
+          msg_interceptor_ctx
          }).
 
 -record(pending_ack, {
@@ -509,7 +510,11 @@ init([Channel, ReaderPid, WriterPid, ConnPid, ConnName, Protocol, User, VHost,
                             consumer_timeout = ConsumerTimeout,
                             authz_context = OptionalVariables,
                             max_consumers = MaxConsumers,
-                            writer_gc_threshold = GCThreshold
+                            writer_gc_threshold = GCThreshold,
+                            msg_interceptor_ctx = #{protocol => amqp091,
+                                                    username => User#user.username,
+                                                    vhost => VHost,
+                                                    conn_name => ConnName}
                            },
                 limiter = Limiter,
                 tx                      = none,
@@ -813,6 +818,7 @@ get_consumer_timeout() ->
         _ ->
             undefined
     end.
+
 %%---------------------------------------------------------------------------
 
 reply(Reply, NewState) -> {reply, Reply, next_state(NewState), hibernate}.
@@ -1167,7 +1173,8 @@ handle_method(#'basic.publish'{exchange    = ExchangeNameBin,
                                                 user = #user{username = Username} = User,
                                                 trace_state = TraceState,
                                                 authz_context = AuthzContext,
-                                                writer_gc_threshold = GCThreshold
+                                                writer_gc_threshold = GCThreshold,
+                                                msg_interceptor_ctx = MsgInterceptorCtx
                                                },
                                    tx               = Tx,
                                    confirm_enabled  = ConfirmEnabled,
@@ -1206,7 +1213,9 @@ 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_message_interceptor:intercept(Message0,
+                                                           MsgInterceptorCtx,
+                                                           incoming_message_interceptors),
             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_message_interceptor.erl

@@ -1,50 +1,49 @@
-%% This Source Code Form is subject to the terms of the Mozilla Public
-%% License, v. 2.0. If a copy of the MPL was not distributed with this
-%% file, You can obtain one at https://mozilla.org/MPL/2.0/.
-%%
-%% Copyright (c) 2007-2025 Broadcom. All Rights Reserved. The term “Broadcom” refers to Broadcom Inc. and/or its subsidiaries. All rights reserved.
-
-%% This module exists since 3.12 replacing plugins rabbitmq-message-timestamp
-%% and rabbitmq-routing-node-stamp. Instead of using these plugins, RabbitMQ core can
-%% now be configured to add such headers. This enables non-AMQP 0.9.1 protocols (that
-%% do not use rabbit_channel) to also add AMQP 0.9.1 headers to incoming messages.
 -module(rabbit_message_interceptor).
--include("mc.hrl").
 
--export([intercept/1]).
+-export([intercept/3,
+         set_msg_annotation/4]).
+
+-type protocol() :: amqp091 | amqp10 | mqtt310 | mqtt311 | mqtt50.
 
--define(HEADER_TIMESTAMP, <<"timestamp_in_ms">>).
--define(HEADER_ROUTING_NODE, <<"x-routed-by">>).
+-type msg_interceptor_ctx() :: #{protocol := protocol(),
+                                 vhost := binary(),
+                                 username := binary(),
+                                 conn_name => binary(),
+                                 atom() => term()}.
 
--spec intercept(mc:state()) -> mc:state().
-intercept(Msg) ->
-    Interceptors = persistent_term:get(incoming_message_interceptors, []),
-    lists:foldl(fun({InterceptorName, Overwrite}, M) ->
-                        intercept(M, InterceptorName, Overwrite)
-                end, Msg, Interceptors).
+-callback intercept(Msg, MsgInterceptorCtx, Group, Config) -> Msg when
+    Msg :: mc:state(),
+    MsgInterceptorCtx :: msg_interceptor_ctx(),
+    Group :: incoming_message_interceptors | outgoing_message_interceptors,
+    Config :: #{atom() := term()}.
 
-intercept(Msg, set_header_routing_node, Overwrite) ->
-    Node = atom_to_binary(node()),
-    set_annotation(Msg, ?HEADER_ROUTING_NODE, Node, Overwrite);
-intercept(Msg0, set_header_timestamp, Overwrite) ->
-    Ts = mc:get_annotation(?ANN_RECEIVED_AT_TIMESTAMP, Msg0),
-    Msg = set_annotation(Msg0, ?HEADER_TIMESTAMP, Ts, Overwrite),
-    set_timestamp(Msg, Ts, Overwrite).
+-spec intercept(Msg, MsgInterceptorCtx, Group) -> Msg when
+    Msg :: mc:state(),
+    MsgInterceptorCtx :: map(),
+    Group :: incoming_message_interceptors | outgoing_message_interceptors.
+intercept(Msg, MsgInterceptorCtx, Group) ->
+    Interceptors = persistent_term:get(Group, []),
+    lists:foldl(fun({Module, Config}, Msg0) ->
+                        try
+                            Module:intercept(Msg0,
+                                             MsgInterceptorCtx,
+                                             Group,
+                                             Config)
+                        catch
+                            error:undef ->
+                                Msg0
+                        end
+                end, Msg , Interceptors).
 
--spec set_annotation(mc:state(), mc:ann_key(), mc:ann_value(), boolean()) -> mc:state().
-set_annotation(Msg, Key, Value, Overwrite) ->
+-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.
-
--spec set_timestamp(mc:state(), pos_integer(), boolean()) -> mc:state().
-set_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_message_interceptor_routing_node.erl

@@ -0,0 +1,14 @@
+-module(rabbit_message_interceptor_routing_node).
+-behaviour(rabbit_message_interceptor).
+
+-define(HEADER_ROUTING_NODE, <<"x-routed-by">>).
+
+-export([intercept/4]).
+
+intercept(Msg, _MsgInterceptorCtx, _Group, Config) ->
+    Node = atom_to_binary(node()),
+    Overwrite = maps:get(overwrite, Config, false),
+    rabbit_message_interceptor:set_msg_annotation(Msg,
+                                                  ?HEADER_ROUTING_NODE,
+                                                  Node,
+                                                  Overwrite).

deps/rabbit/src/rabbit_message_interceptor_timestamp.erl

@@ -0,0 +1,26 @@
+-module(rabbit_message_interceptor_timestamp).
+-behaviour(rabbit_message_interceptor).
+
+-include("mc.hrl").
+
+-define(HEADER_TIMESTAMP, <<"timestamp_in_ms">>).
+
+-export([intercept/4]).
+
+intercept(Msg0, _MsgInterceptorCtx, _Group, Config) ->
+    Ts = mc:get_annotation(?ANN_RECEIVED_AT_TIMESTAMP, Msg0),
+    Overwrite = maps:get(overwrite, Config, false),
+    Msg = rabbit_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/test/amqp_client_SUITE.erl

@@ -4380,8 +4380,11 @@ available_messages(QType, Config) ->
 
 incoming_message_interceptors(Config) ->
     Key = ?FUNCTION_NAME,
-    ok = rpc(Config, persistent_term, put, [Key, [{set_header_routing_node, false},
-                                                  {set_header_timestamp, false}]]),
+    ok = rpc(Config,
+             persistent_term,
+             put,
+             [Key, [{rabbit_message_interceptor_routing_node, #{overwrite => false}},
+                    {rabbit_message_interceptor_timestamp, #{overwrite => false}}]]),
     Stream = <<"my stream">>,
     QQName = <<"my quorum queue">>,
     {_, Session, LinkPair} = Init = init(Config),