Support outgoing message interceptors

Author: ansd

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_msg_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_msg_interceptor_SUITE
 ,   rabbit_stream_coordinator_SUITE
 ,   rabbit_stream_sac_coordinator_SUITE
 ,   rabbitmq_4_0_deprecations_SUITE

deps/rabbit/priv/schema/rabbit.schema

@@ -2664,27 +2664,52 @@ end}.
 %%
 %% Message interceptors
 %%
-{mapping, "message_interceptors.incoming.$interceptor.overwrite", "rabbit.incoming_message_interceptors", [
+
+{mapping, "message_interceptors.$stage.$name.$key", "rabbit.message_interceptors", [
     {datatype, {enum, [true, false]}}]}.
 
-{translation, "rabbit.incoming_message_interceptors",
+{translation, "rabbit.message_interceptors",
  fun(Conf) ->
          case cuttlefish_variable:filter_by_prefix("message_interceptors", Conf) of
              [] ->
                  cuttlefish:unset();
              L ->
-                 [begin
-                      Interceptor = list_to_atom(Interceptor0),
-                      Mod = case Interceptor of
-                                set_header_timestamp ->
-                                    rabbit_message_interceptor_timestamp;
-                                set_header_routing_node ->
-                                    rabbit_message_interceptor_routing_node;
-                                _ ->
-                                    cuttlefish:invalid(io_lib:format("~p is invalid", [Interceptor]))
-                            end,
-                      {Mod, #{overwrite => Overwrite}}
-                  end || {["message_interceptors", "incoming", Interceptor0, "overwrite"], Overwrite} <- L]
+                 lists:foldr(
+                   fun({["message_interceptors", "incoming", "set_header_routing_node", "overwrite"], Overwrite}, Acc)
+                         when is_boolean(Overwrite) ->
+                           Mod = rabbit_msg_interceptor_routing_node,
+                           Cfg = #{overwrite => Overwrite},
+                           [{Mod, Cfg} | Acc];
+                      ({["message_interceptors", "incoming", "set_header_timestamp", "overwrite"], Overwrite}, Acc)
+                        when is_boolean(Overwrite) ->
+                           Mod = rabbit_msg_interceptor_timestamp,
+                           Cfg = #{incoming => true,
+                                   overwrite => Overwrite},
+                           case lists:keytake(Mod, 1, Acc) of
+                               false ->
+                                   [{Mod, Cfg} | Acc];
+                               {value, {Mod, Cfg1}, Acc1} ->
+                                   Cfg2 = maps:merge(Cfg1, Cfg),
+                                   [{Mod, Cfg2} | Acc1]
+                           end;
+                      ({["message_interceptors", "outgoing", "timestamp", "enabled"], Enabled}, Acc) ->
+                           case Enabled of
+                               true ->
+                                   Mod = rabbit_msg_interceptor_timestamp,
+                                   Cfg = #{outgoing => true},
+                                   case lists:keytake(Mod, 1, Acc) of
+                                       false ->
+                                           [{Mod, Cfg} | Acc];
+                                       {value, {Mod, Cfg1}, Acc1} ->
+                                           Cfg2 = maps:merge(Cfg1, Cfg),
+                                           [{Mod, Cfg2} | Acc1]
+                                   end;
+                               false ->
+                                   Acc
+                           end;
+                      (Other, _Acc) ->
+                           cuttlefish:invalid(io_lib:format("~p is invalid", [Other]))
+                   end, [], L)
          end
  end
 }.

deps/rabbit/src/mc_amqpl.erl

@@ -462,7 +462,6 @@ protocol_state(#content{properties = #'P_basic'{headers = H00,
                                                 priority = Priority0,
                                                 delivery_mode = DeliveryMode0} = B0} = C,
                Anns) ->
-    %% Add any x- annotations as headers
     H0 = case H00 of
              undefined -> [];
              _ ->
@@ -474,6 +473,7 @@ protocol_state(#content{properties = #'P_basic'{headers = H00,
                    _ ->
                        H0
                end,
+    %% Add any x- annotations as headers
     Headers1 = maps:fold(
                  fun (<<"x-", _/binary>> = Key, Val, H) when is_integer(Val) ->
                          [{Key, long, Val} | H];

deps/rabbit/src/rabbit.erl

@@ -1658,8 +1658,8 @@ persist_static_configuration() ->
        classic_queue_store_v2_check_crc32
       ]),
 
-    Interceptors = application:get_env(?MODULE, incoming_message_interceptors, []),
-    ok = rabbit_message_interceptor:add(Interceptors, incoming_message_interceptors),
+    Interceptors = application:get_env(?MODULE, message_interceptors, []),
+    ok = rabbit_msg_interceptor:add(Interceptors),
 
     %% Disallow the following two cases:
     %% 1. Negative values

deps/rabbit/src/rabbit_amqp_session.erl

@@ -284,7 +284,7 @@
           max_incoming_window :: pos_integer(),
           max_link_credit :: pos_integer(),
           max_queue_credit :: pos_integer(),
-          msg_interceptor_ctx :: rabbit_message_interceptor:context()
+          msg_interceptor_ctx :: rabbit_msg_interceptor:context()
          }).
 
 -record(state, {
@@ -2164,7 +2164,8 @@ handle_deliver(ConsumerTag, AckRequired,
                                          conn_name = ConnName,
                                          channel_num = ChannelNum,
                                          user = #user{username = Username},
-                                         trace_state = Trace}}) ->
+                                         trace_state = Trace,
+                                         msg_interceptor_ctx = MsgIcptCtx}}) ->
     Handle = ctag_to_handle(ConsumerTag),
     case OutgoingLinks0 of
         #{Handle := #outgoing_link{queue_type = QType,
@@ -2180,7 +2181,8 @@ handle_deliver(ConsumerTag, AckRequired,
                           message_format = ?UINT(?MESSAGE_FORMAT),
                           settled = SendSettled},
             Mc1 = mc:convert(mc_amqp, Mc0),
-            Mc = mc:set_annotation(redelivered, Redelivered, Mc1),
+            Mc2 = mc:set_annotation(redelivered, Redelivered, Mc1),
+            Mc = rabbit_msg_interceptor:intercept_outgoing(Mc2, MsgIcptCtx),
             Sections = mc:protocol_state(Mc),
             validate_message_size(Sections, MaxMessageSize),
             Frames = transfer_frames(Transfer, Sections, MaxFrameSize),
@@ -2417,7 +2419,7 @@ incoming_link_transfer(
                              conn_name = ConnName,
                              channel_num = ChannelNum,
                              max_link_credit = MaxLinkCredit,
-                             msg_interceptor_ctx = MsgInterceptorCtx}}) ->
+                             msg_interceptor_ctx = MsgIcptCtx}}) ->
 
     {PayloadBin, DeliveryId, Settled} =
     case MultiTransfer of
@@ -2445,9 +2447,7 @@ incoming_link_transfer(
             check_user_id(Mc1, User),
             TopicPermCache = check_write_permitted_on_topics(
                                X, User, RoutingKeys, TopicPermCache0),
-            Mc2 = rabbit_message_interceptor:intercept(Mc1,
-                                                       MsgInterceptorCtx,
-                                                       incoming_message_interceptors),
+            Mc2 = rabbit_msg_interceptor:intercept_incoming(Mc1, MsgIcptCtx),
             QNames = rabbit_exchange:route(X, Mc2, #{return_binding_keys => true}),
             rabbit_trace:tap_in(Mc2, QNames, ConnName, ChannelNum, Username, Trace),
             Opts = #{correlation => {HandleInt, DeliveryId}},

deps/rabbit/src/rabbit_channel.erl

@@ -111,7 +111,7 @@
           max_consumers,  % taken from rabbit.consumer_max_per_channel
           %% defines how ofter gc will be executed
           writer_gc_threshold,
-          msg_interceptor_ctx :: rabbit_message_interceptor:context()
+          msg_interceptor_ctx :: rabbit_msg_interceptor:context()
          }).
 
 -record(pending_ack, {
@@ -662,13 +662,14 @@ handle_cast({deliver_reply, _K, _Del},
     noreply(State);
 handle_cast({deliver_reply, _K, _Msg}, State = #ch{reply_consumer = none}) ->
     noreply(State);
-handle_cast({deliver_reply, Key, Msg},
-            State = #ch{cfg = #conf{writer_pid = WriterPid},
+handle_cast({deliver_reply, Key, Mc},
+            State = #ch{cfg = #conf{writer_pid = WriterPid,
+                                    msg_interceptor_ctx = MsgIcptCtx},
                         next_tag = DeliveryTag,
                         reply_consumer = {ConsumerTag, _Suffix, Key}}) ->
-    Content = mc:protocol_state(mc:convert(mc_amqpl, Msg)),
-    ExchName = mc:exchange(Msg),
-    [RoutingKey | _] = mc:routing_keys(Msg),
+    ExchName = mc:exchange(Mc),
+    [RoutingKey | _] = mc:routing_keys(Mc),
+    Content = outgoing_content(Mc, MsgIcptCtx),
     ok = rabbit_writer:send_command(
            WriterPid,
            #'basic.deliver'{consumer_tag = ConsumerTag,
@@ -1174,7 +1175,7 @@ handle_method(#'basic.publish'{exchange    = ExchangeNameBin,
                                                 trace_state = TraceState,
                                                 authz_context = AuthzContext,
                                                 writer_gc_threshold = GCThreshold,
-                                                msg_interceptor_ctx = MsgInterceptorCtx
+                                                msg_interceptor_ctx = MsgIcptCtx
                                                },
                                    tx               = Tx,
                                    confirm_enabled  = ConfirmEnabled,
@@ -1214,9 +1215,7 @@ handle_method(#'basic.publish'{exchange    = ExchangeNameBin,
         {ok, Message0} ->
             check_write_permitted_on_topics(Exchange, User, Message0, AuthzContext),
             check_user_id_header(Message0, User),
-            Message = rabbit_message_interceptor:intercept(Message0,
-                                                           MsgInterceptorCtx,
-                                                           incoming_message_interceptors),
+            Message = rabbit_msg_interceptor:intercept_incoming(Message0, MsgIcptCtx),
             QNames = rabbit_exchange:route(Exchange, Message, #{return_binding_keys => true}),
             [deliver_reply(RK, Message) || {virtual_reply_queue, RK} <- QNames],
             Queues = rabbit_amqqueue:lookup_many(QNames),
@@ -2601,15 +2600,15 @@ handle_deliver(CTag, Ack, Msgs, State) when is_list(Msgs) ->
                 end, State, Msgs).
 
 handle_deliver0(ConsumerTag, AckRequired,
-                {QName, QPid, _MsgId, Redelivered, MsgCont0} = Msg,
+                {QName, QPid, _MsgId, Redelivered, Mc} = Msg,
                State = #ch{cfg = #conf{writer_pid = WriterPid,
-                                       writer_gc_threshold = GCThreshold},
+                                       writer_gc_threshold = GCThreshold,
+                                       msg_interceptor_ctx = MsgIcptCtx},
                            next_tag   = DeliveryTag,
                            queue_states = Qs}) ->
-    Exchange = mc:exchange(MsgCont0),
-    [RoutingKey | _] = mc:routing_keys(MsgCont0),
-    MsgCont = mc:convert(mc_amqpl, MsgCont0),
-    Content = mc:protocol_state(MsgCont),
+    Exchange = mc:exchange(Mc),
+    [RoutingKey | _] = mc:routing_keys(Mc),
+    Content = outgoing_content(Mc, MsgIcptCtx),
     Deliver = #'basic.deliver'{consumer_tag = ConsumerTag,
                                delivery_tag = DeliveryTag,
                                redelivered  = Redelivered,
@@ -2630,12 +2629,11 @@ handle_deliver0(ConsumerTag, AckRequired,
     record_sent(deliver, QueueType, ConsumerTag, AckRequired, Msg, State).
 
 handle_basic_get(WriterPid, DeliveryTag, NoAck, MessageCount,
-                 Msg0 = {_QName, _QPid, _MsgId, Redelivered, MsgCont0},
+                 Msg0 = {_QName, _QPid, _MsgId, Redelivered, Mc},
                  QueueType, State) ->
-    Exchange = mc:exchange(MsgCont0),
-    [RoutingKey | _] = mc:routing_keys(MsgCont0),
-    MsgCont = mc:convert(mc_amqpl, MsgCont0),
-    Content = mc:protocol_state(MsgCont),
+    Exchange = mc:exchange(Mc),
+    [RoutingKey | _] = mc:routing_keys(Mc),
+    Content = outgoing_content(Mc, State#ch.cfg#conf.msg_interceptor_ctx),
     ok = rabbit_writer:send_command(
            WriterPid,
            #'basic.get_ok'{delivery_tag  = DeliveryTag,
@@ -2646,6 +2644,11 @@ handle_basic_get(WriterPid, DeliveryTag, NoAck, MessageCount,
            Content),
     {noreply, record_sent(get, QueueType, DeliveryTag, not(NoAck), Msg0, State)}.
 
+outgoing_content(Mc, MsgIcptCtx) ->
+    Mc1 = mc:convert(mc_amqpl, Mc),
+    Mc2 = rabbit_msg_interceptor:intercept_outgoing(Mc1, MsgIcptCtx),
+    mc:protocol_state(Mc2).
+
 init_tick_timer(State = #ch{tick_timer = undefined}) ->
     {ok, Interval} = application:get_env(rabbit, channel_tick_interval),
     State#ch{tick_timer = erlang:send_after(Interval, self(), tick)};

deps/rabbit/src/rabbit_message_interceptor_timestamp.erl

@@ -4,12 +4,13 @@
 %%
 %% Copyright (c) 2007-2025 Broadcom. All Rights Reserved. The term “Broadcom” refers to Broadcom Inc. and/or its subsidiaries. All rights reserved.
 
--module(rabbit_message_interceptor).
+-module(rabbit_msg_interceptor).
 
 %% client API
--export([intercept/3,
-         add/2,
-         remove/2]).
+-export([intercept_incoming/2,
+         intercept_outgoing/2,
+         add/1,
+         remove/1]).
 %% helpers for behaviour implementations
 -export([set_annotation/4]).
 
@@ -20,46 +21,58 @@
                      username := rabbit_types:username(),
                      connection_name := binary(),
                      atom() => term()}.
--type group() :: incoming_message_interceptors |
-                 outgoing_message_interceptors.
 -type config() :: #{atom() => term()}.
 -type interceptor() :: {module(), config()}.
+-type interceptors() :: [interceptor()].
+-type stage() :: incoming | outgoing.
 
+-define(KEY, message_interceptors).
 
 -export_type([context/0]).
 
--callback intercept(mc:state(), context(), group(), config()) ->
+-callback intercept(mc:state(), context(), stage(), config()) ->
     mc:state().
 
--spec intercept(mc:state(), context(), group()) ->
+-spec intercept_incoming(mc:state(), context()) ->
     mc:state().
-intercept(Msg, Ctx, Group) ->
-    Interceptors = persistent_term:get(Group, []),
+intercept_incoming(Msg, Ctx) ->
+    intercept(Msg, Ctx, incoming).
+
+-spec intercept_outgoing(mc:state(), context()) ->
+    mc:state().
+intercept_outgoing(Msg, Ctx) ->
+    intercept(Msg, Ctx, outgoing).
+
+intercept(Msg, Ctx, Stage) ->
+    Interceptors = persistent_term:get(?KEY),
     lists:foldl(fun({Mod, Config}, Msg0) ->
-                        Mod:intercept(Msg0, Ctx, Group, Config)
+                        Mod:intercept(Msg0, Ctx, Stage, Config)
                 end, Msg, Interceptors).
 
--spec set_annotation(mc:state(), mc:ann_key(), mc:ann_value(), boolean()) ->
+-spec set_annotation(mc:state(), mc:ann_key(), mc:ann_value(),
+                     Overwrite :: boolean()) ->
     mc:state().
-set_annotation(Msg, Key, Value, Overwrite) ->
-    case {mc:x_header(Key, Msg), Overwrite} of
-        {Val, false} when Val =/= undefined ->
-            Msg;
+set_annotation(Msg, Key, Value, true) ->
+    mc:set_annotation(Key, Value, Msg);
+set_annotation(Msg, Key, Value, false) ->
+    case mc:x_header(Key, Msg) of
+        undefined ->
+            mc:set_annotation(Key, Value, Msg);
         _ ->
-            mc:set_annotation(Key, Value, Msg)
+            Msg
     end.
 
--spec add([interceptor()], group()) -> ok.
-add(Interceptors, Group) ->
+-spec add(interceptors()) -> ok.
+add(Interceptors) ->
     %% validation
     lists:foreach(fun({Mod, #{}}) ->
                           case erlang:function_exported(Mod, intercept, 4) of
                               true -> ok;
                               false -> error(Mod)
                           end
                   end, Interceptors),
-    persistent_term:put(Group, persistent_term:get(Group, []) ++ Interceptors).
+    persistent_term:put(?KEY, persistent_term:get(?KEY, []) ++ Interceptors).
 
--spec remove([interceptor()], group()) -> ok.
-remove(Interceptors, Group) ->
-    persistent_term:put(Group, persistent_term:get(Group, []) -- Interceptors).
+-spec remove(interceptors()) -> ok.
+remove(Interceptors) ->
+    persistent_term:put(?KEY, persistent_term:get(?KEY, []) -- Interceptors).