rabbitmq-server/deps/rabbitmq_mqtt/src/rabbit_mqtt_util.erl

212 lines
6.9 KiB
Erlang

%% 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-2023 VMware, Inc. or its affiliates. All rights reserved.
%%
-module(rabbit_mqtt_util).
-include_lib("rabbit_common/include/resource.hrl").
-include("rabbit_mqtt.hrl").
-include("rabbit_mqtt_packet.hrl").
-export([queue_name_bin/2,
qos_from_queue_name/2,
env/1,
table_lookup/2,
path_for/2,
path_for/3,
vhost_name_to_table_name/1,
register_clientid/2,
remove_duplicate_clientid_connections/2,
init_sparkplug/0,
mqtt_to_amqp/1,
amqp_to_mqtt/1,
truncate_binary/2
]).
-define(MAX_TOPIC_TRANSLATION_CACHE_SIZE, 12).
-define(SPARKPLUG_MP_MQTT_TO_AMQP, sparkplug_mp_mqtt_to_amqp).
-define(SPARKPLUG_MP_AMQP_TO_MQTT, sparkplug_mp_amqp_to_mqtt).
-spec queue_name_bin(binary(), qos()) ->
binary().
queue_name_bin(ClientId, QoS) ->
Prefix = queue_name_prefix(ClientId),
queue_name0(Prefix, QoS).
queue_name0(Prefix, ?QOS_0) ->
<<Prefix/binary, "0">>;
queue_name0(Prefix, ?QOS_1) ->
<<Prefix/binary, "1">>.
-spec qos_from_queue_name(rabbit_amqqueue:name(), binary()) ->
qos() | no_consuming_queue.
qos_from_queue_name(#resource{name = Name}, ClientId) ->
Prefix = queue_name_prefix(ClientId),
PrefixSize = erlang:byte_size(Prefix),
case Name of
<<Prefix:PrefixSize/binary, "0">> ->
?QOS_0;
<<Prefix:PrefixSize/binary, "1">> ->
?QOS_1;
_ ->
no_consuming_queue
end.
queue_name_prefix(ClientId) ->
<<"mqtt-subscription-", ClientId/binary, "qos">>.
-spec init_sparkplug() -> ok.
init_sparkplug() ->
case env(sparkplug) of
true ->
{ok, M2A_SpRe} = re:compile("^sp[AB]v\\d+\\.\\d+/"),
{ok, A2M_SpRe} = re:compile("^sp[AB]v\\d+___\\d+\\."),
ok = persistent_term:put(?SPARKPLUG_MP_MQTT_TO_AMQP, M2A_SpRe),
ok = persistent_term:put(?SPARKPLUG_MP_AMQP_TO_MQTT, A2M_SpRe);
_ ->
ok
end.
-spec mqtt_to_amqp(binary()) -> binary().
mqtt_to_amqp(Topic) ->
T = case persistent_term:get(?SPARKPLUG_MP_MQTT_TO_AMQP, no_sparkplug) of
no_sparkplug ->
Topic;
M2A_SpRe ->
case re:run(Topic, M2A_SpRe) of
nomatch ->
Topic;
{match, _} ->
string:replace(Topic, ".", "___", leading)
end
end,
cached(mta_cache, fun to_amqp/1, T).
-spec amqp_to_mqtt(binary()) -> binary().
amqp_to_mqtt(Topic) ->
T = cached(atm_cache, fun to_mqtt/1, Topic),
case persistent_term:get(?SPARKPLUG_MP_AMQP_TO_MQTT, no_sparkplug) of
no_sparkplug ->
T;
A2M_SpRe ->
case re:run(Topic, A2M_SpRe) of
nomatch ->
T;
{match, _} ->
T1 = string:replace(T, "___", ".", leading),
erlang:iolist_to_binary(T1)
end
end.
cached(CacheName, Fun, Arg) ->
Cache = case get(CacheName) of
undefined ->
[];
Other ->
Other
end,
case lists:keyfind(Arg, 1, Cache) of
{_, V} ->
V;
false ->
V = Fun(Arg),
CacheTail = lists:sublist(Cache, ?MAX_TOPIC_TRANSLATION_CACHE_SIZE - 1),
put(CacheName, [{Arg, V} | CacheTail]),
V
end.
%% amqp mqtt descr
%% * + match one topic level
%% # # match multiple topic levels
%% . / topic level separator
to_amqp(T0) ->
T1 = string:replace(T0, "/", ".", all),
T2 = string:replace(T1, "+", "*", all),
erlang:iolist_to_binary(T2).
to_mqtt(T0) ->
T1 = string:replace(T0, "*", "+", all),
T2 = string:replace(T1, ".", "/", all),
erlang:iolist_to_binary(T2).
-spec env(atom()) -> any().
env(Key) ->
case application:get_env(?APP_NAME, Key) of
{ok, Val} -> coerce_env_value(Key, Val);
undefined -> undefined
end.
coerce_env_value(default_pass, Val) -> rabbit_data_coercion:to_binary(Val);
coerce_env_value(default_user, Val) -> rabbit_data_coercion:to_binary(Val);
coerce_env_value(exchange, Val) -> rabbit_data_coercion:to_binary(Val);
coerce_env_value(vhost, Val) -> rabbit_data_coercion:to_binary(Val);
coerce_env_value(_, Val) -> Val.
-spec table_lookup(rabbit_framing:amqp_table() | undefined, binary()) ->
tuple() | undefined.
table_lookup(undefined, _Key) ->
undefined;
table_lookup(Table, Key) ->
rabbit_misc:table_lookup(Table, Key).
vhost_name_to_dir_name(VHost) ->
vhost_name_to_dir_name(VHost, ".ets").
vhost_name_to_dir_name(VHost, Suffix) ->
<<Num:128>> = erlang:md5(VHost),
"mqtt_retained_" ++ rabbit_misc:format("~36.16.0b", [Num]) ++ Suffix.
-spec path_for(file:name_all(), rabbit_types:vhost()) -> file:filename_all().
path_for(Dir, VHost) ->
filename:join(Dir, vhost_name_to_dir_name(VHost)).
-spec path_for(file:name_all(), rabbit_types:vhost(), string()) -> file:filename_all().
path_for(Dir, VHost, Suffix) ->
filename:join(Dir, vhost_name_to_dir_name(VHost, Suffix)).
-spec vhost_name_to_table_name(rabbit_types:vhost()) ->
atom().
vhost_name_to_table_name(VHost) ->
<<Num:128>> = erlang:md5(VHost),
list_to_atom("rabbit_mqtt_retained_" ++ rabbit_misc:format("~36.16.0b", [Num])).
-spec register_clientid(rabbit_types:vhost(), binary()) -> ok.
register_clientid(Vhost, ClientId)
when is_binary(Vhost), is_binary(ClientId) ->
PgGroup = {Vhost, ClientId},
ok = pg:join(persistent_term:get(?PG_SCOPE), PgGroup, self()),
case rabbit_mqtt_ff:track_client_id_in_ra() of
true ->
%% Ra node takes care of removing duplicate client ID connections.
ok;
false ->
ok = erpc:multicast([node() | nodes()],
?MODULE,
remove_duplicate_clientid_connections,
[PgGroup, self()])
end.
-spec remove_duplicate_clientid_connections({rabbit_types:vhost(), binary()}, pid()) -> ok.
remove_duplicate_clientid_connections(PgGroup, PidToKeep) ->
try persistent_term:get(?PG_SCOPE) of
PgScope ->
Pids = pg:get_local_members(PgScope, PgGroup),
lists:foreach(fun(Pid) ->
gen_server:cast(Pid, duplicate_id)
end, Pids -- [PidToKeep])
catch _:badarg ->
%% MQTT supervision tree on this node not fully started
ok
end.
-spec truncate_binary(binary(), non_neg_integer()) -> binary().
truncate_binary(Bin, Size)
when is_binary(Bin) andalso byte_size(Bin) =< Size ->
Bin;
truncate_binary(Bin, Size)
when is_binary(Bin) ->
binary:part(Bin, 0, Size).