469 lines
17 KiB
Erlang
469 lines
17 KiB
Erlang
%%--------------------------------------------------------------------
|
|
%% Copyright (c) 2020-2021 EMQ Technologies Co., Ltd. All Rights Reserved.
|
|
%%
|
|
%% Licensed under the Apache License, Version 2.0 (the "License");
|
|
%% you may not use this file except in compliance with the License.
|
|
%% You may obtain a copy of the License at
|
|
%%
|
|
%% http://www.apache.org/licenses/LICENSE-2.0
|
|
%%
|
|
%% Unless required by applicable law or agreed to in writing, software
|
|
%% distributed under the License is distributed on an "AS IS" BASIS,
|
|
%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
%% See the License for the specific language governing permissions and
|
|
%% limitations under the License.
|
|
%%--------------------------------------------------------------------
|
|
|
|
%% @doc Stomp Protocol Processor.
|
|
-module(emqx_stomp_protocol).
|
|
|
|
-include("emqx_stomp.hrl").
|
|
|
|
-include_lib("emqx/include/emqx.hrl").
|
|
-include_lib("emqx/include/logger.hrl").
|
|
-include_lib("emqx/include/emqx_mqtt.hrl").
|
|
|
|
-logger_header("[Stomp-Proto]").
|
|
|
|
-import(proplists, [get_value/2, get_value/3]).
|
|
|
|
%% API
|
|
-export([ init/2
|
|
, info/1
|
|
]).
|
|
|
|
-export([ received/2
|
|
, send/2
|
|
, shutdown/2
|
|
, timeout/3
|
|
]).
|
|
|
|
%% for trans callback
|
|
-export([ handle_recv_send_frame/2
|
|
, handle_recv_ack_frame/2
|
|
, handle_recv_nack_frame/2
|
|
]).
|
|
|
|
-record(pstate, {
|
|
peername,
|
|
heartfun,
|
|
sendfun,
|
|
connected = false,
|
|
proto_ver,
|
|
proto_name,
|
|
heart_beats,
|
|
login,
|
|
allow_anonymous,
|
|
default_user,
|
|
subscriptions = [],
|
|
timers :: #{atom() => disable | undefined | reference()},
|
|
transaction :: #{binary() => list()}
|
|
}).
|
|
|
|
-define(TIMER_TABLE, #{
|
|
incoming_timer => incoming,
|
|
outgoing_timer => outgoing,
|
|
clean_trans_timer => clean_trans
|
|
}).
|
|
|
|
-define(TRANS_TIMEOUT, 60000).
|
|
|
|
-type(pstate() :: #pstate{}).
|
|
|
|
%% @doc Init protocol
|
|
init(#{peername := Peername,
|
|
sendfun := SendFun,
|
|
heartfun := HeartFun}, Env) ->
|
|
AllowAnonymous = get_value(allow_anonymous, Env, false),
|
|
DefaultUser = get_value(default_user, Env),
|
|
#pstate{peername = Peername,
|
|
heartfun = HeartFun,
|
|
sendfun = SendFun,
|
|
timers = #{},
|
|
transaction = #{},
|
|
allow_anonymous = AllowAnonymous,
|
|
default_user = DefaultUser}.
|
|
|
|
info(#pstate{connected = Connected,
|
|
proto_ver = ProtoVer,
|
|
proto_name = ProtoName,
|
|
heart_beats = Heartbeats,
|
|
login = Login,
|
|
subscriptions = Subscriptions}) ->
|
|
[{connected, Connected},
|
|
{proto_ver, ProtoVer},
|
|
{proto_name, ProtoName},
|
|
{heart_beats, Heartbeats},
|
|
{login, Login},
|
|
{subscriptions, Subscriptions}].
|
|
|
|
-spec(received(stomp_frame(), pstate())
|
|
-> {ok, pstate()}
|
|
| {error, any(), pstate()}
|
|
| {stop, any(), pstate()}).
|
|
received(Frame = #stomp_frame{command = <<"STOMP">>}, State) ->
|
|
received(Frame#stomp_frame{command = <<"CONNECT">>}, State);
|
|
|
|
received(#stomp_frame{command = <<"CONNECT">>, headers = Headers},
|
|
State = #pstate{connected = false, allow_anonymous = AllowAnonymous, default_user = DefaultUser}) ->
|
|
case negotiate_version(header(<<"accept-version">>, Headers)) of
|
|
{ok, Version} ->
|
|
Login = header(<<"login">>, Headers),
|
|
Passc = header(<<"passcode">>, Headers),
|
|
case check_login(Login, Passc, AllowAnonymous, DefaultUser) of
|
|
true ->
|
|
emqx_logger:set_metadata_clientid(Login),
|
|
|
|
Heartbeats = parse_heartbeats(header(<<"heart-beat">>, Headers, <<"0,0">>)),
|
|
NState = start_heartbeart_timer(Heartbeats, State#pstate{connected = true,
|
|
proto_ver = Version, login = Login}),
|
|
send(connected_frame([{<<"version">>, Version},
|
|
{<<"heart-beat">>, reverse_heartbeats(Heartbeats)}]), NState);
|
|
false ->
|
|
_ = send(error_frame(undefined, <<"Login or passcode error!">>), State),
|
|
{error, login_or_passcode_error, State}
|
|
end;
|
|
{error, Msg} ->
|
|
_ = send(error_frame([{<<"version">>, <<"1.0,1.1,1.2">>},
|
|
{<<"content-type">>, <<"text/plain">>}], undefined, Msg), State),
|
|
{error, unsupported_version, State}
|
|
end;
|
|
|
|
received(#stomp_frame{command = <<"CONNECT">>}, State = #pstate{connected = true}) ->
|
|
{error, unexpected_connect, State};
|
|
|
|
received(Frame = #stomp_frame{command = <<"SEND">>, headers = Headers}, State) ->
|
|
case header(<<"transaction">>, Headers) of
|
|
undefined -> {ok, handle_recv_send_frame(Frame, State)};
|
|
TransactionId -> add_action(TransactionId, {fun ?MODULE:handle_recv_send_frame/2, [Frame]}, receipt_id(Headers), State)
|
|
end;
|
|
|
|
received(#stomp_frame{command = <<"SUBSCRIBE">>, headers = Headers},
|
|
State = #pstate{subscriptions = Subscriptions}) ->
|
|
Id = header(<<"id">>, Headers),
|
|
Topic = header(<<"destination">>, Headers),
|
|
Ack = header(<<"ack">>, Headers, <<"auto">>),
|
|
{ok, State1} = case lists:keyfind(Id, 1, Subscriptions) of
|
|
{Id, Topic, Ack} ->
|
|
{ok, State};
|
|
false ->
|
|
emqx_broker:subscribe(Topic),
|
|
{ok, State#pstate{subscriptions = [{Id, Topic, Ack}|Subscriptions]}}
|
|
end,
|
|
maybe_send_receipt(receipt_id(Headers), State1);
|
|
|
|
received(#stomp_frame{command = <<"UNSUBSCRIBE">>, headers = Headers},
|
|
State = #pstate{subscriptions = Subscriptions}) ->
|
|
Id = header(<<"id">>, Headers),
|
|
|
|
{ok, State1} = case lists:keyfind(Id, 1, Subscriptions) of
|
|
{Id, Topic, _Ack} ->
|
|
ok = emqx_broker:unsubscribe(Topic),
|
|
{ok, State#pstate{subscriptions = lists:keydelete(Id, 1, Subscriptions)}};
|
|
false ->
|
|
{ok, State}
|
|
end,
|
|
maybe_send_receipt(receipt_id(Headers), State1);
|
|
|
|
%% ACK
|
|
%% id:12345
|
|
%% transaction:tx1
|
|
%%
|
|
%% ^@
|
|
received(Frame = #stomp_frame{command = <<"ACK">>, headers = Headers}, State) ->
|
|
case header(<<"transaction">>, Headers) of
|
|
undefined -> {ok, handle_recv_ack_frame(Frame, State)};
|
|
TransactionId -> add_action(TransactionId, {fun ?MODULE:handle_recv_ack_frame/2, [Frame]}, receipt_id(Headers), State)
|
|
end;
|
|
|
|
%% NACK
|
|
%% id:12345
|
|
%% transaction:tx1
|
|
%%
|
|
%% ^@
|
|
received(Frame = #stomp_frame{command = <<"NACK">>, headers = Headers}, State) ->
|
|
case header(<<"transaction">>, Headers) of
|
|
undefined -> {ok, handle_recv_nack_frame(Frame, State)};
|
|
TransactionId -> add_action(TransactionId, {fun ?MODULE:handle_recv_nack_frame/2, [Frame]}, receipt_id(Headers), State)
|
|
end;
|
|
|
|
%% BEGIN
|
|
%% transaction:tx1
|
|
%%
|
|
%% ^@
|
|
received(#stomp_frame{command = <<"BEGIN">>, headers = Headers},
|
|
State = #pstate{transaction = Trans}) ->
|
|
Id = header(<<"transaction">>, Headers),
|
|
case maps:get(Id, Trans, undefined) of
|
|
undefined ->
|
|
Ts = erlang:system_time(millisecond),
|
|
NState = ensure_clean_trans_timer(State#pstate{transaction = Trans#{Id => {Ts, []}}}),
|
|
maybe_send_receipt(receipt_id(Headers), NState);
|
|
_ ->
|
|
send(error_frame(receipt_id(Headers), ["Transaction ", Id, " already started"]), State)
|
|
end;
|
|
|
|
%% COMMIT
|
|
%% transaction:tx1
|
|
%%
|
|
%% ^@
|
|
received(#stomp_frame{command = <<"COMMIT">>, headers = Headers},
|
|
State = #pstate{transaction = Trans}) ->
|
|
Id = header(<<"transaction">>, Headers),
|
|
case maps:get(Id, Trans, undefined) of
|
|
{_, Actions} ->
|
|
NState = lists:foldr(fun({Func, Args}, S) ->
|
|
erlang:apply(Func, Args ++ [S])
|
|
end, State#pstate{transaction = maps:remove(Id, Trans)}, Actions),
|
|
maybe_send_receipt(receipt_id(Headers), NState);
|
|
_ ->
|
|
send(error_frame(receipt_id(Headers), ["Transaction ", Id, " not found"]), State)
|
|
end;
|
|
|
|
%% ABORT
|
|
%% transaction:tx1
|
|
%%
|
|
%% ^@
|
|
received(#stomp_frame{command = <<"ABORT">>, headers = Headers},
|
|
State = #pstate{transaction = Trans}) ->
|
|
Id = header(<<"transaction">>, Headers),
|
|
case maps:get(Id, Trans, undefined) of
|
|
{_, _Actions} ->
|
|
NState = State#pstate{transaction = maps:remove(Id, Trans)},
|
|
maybe_send_receipt(receipt_id(Headers), NState);
|
|
_ ->
|
|
send(error_frame(receipt_id(Headers), ["Transaction ", Id, " not found"]), State)
|
|
end;
|
|
|
|
received(#stomp_frame{command = <<"DISCONNECT">>, headers = Headers}, State) ->
|
|
_ = maybe_send_receipt(receipt_id(Headers), State),
|
|
{stop, normal, State}.
|
|
|
|
send(Msg = #message{topic = Topic, headers = Headers, payload = Payload},
|
|
State = #pstate{subscriptions = Subscriptions}) ->
|
|
case lists:keyfind(Topic, 2, Subscriptions) of
|
|
{Id, Topic, Ack} ->
|
|
Headers0 = [{<<"subscription">>, Id},
|
|
{<<"message-id">>, next_msgid()},
|
|
{<<"destination">>, Topic},
|
|
{<<"content-type">>, <<"text/plain">>}],
|
|
Headers1 = case Ack of
|
|
_ when Ack =:= <<"client">> orelse Ack =:= <<"client-individual">> ->
|
|
Headers0 ++ [{<<"ack">>, next_ackid()}];
|
|
_ ->
|
|
Headers0
|
|
end,
|
|
Frame = #stomp_frame{command = <<"MESSAGE">>,
|
|
headers = Headers1 ++ maps:get(stomp_headers, Headers, []),
|
|
body = Payload},
|
|
send(Frame, State);
|
|
false ->
|
|
?LOG(error, "Stomp dropped: ~p", [Msg]),
|
|
{error, dropped, State}
|
|
end;
|
|
|
|
send(Frame, State = #pstate{sendfun = {Fun, Args}}) ->
|
|
?LOG(info, "SEND Frame: ~s", [emqx_stomp_frame:format(Frame)]),
|
|
Data = emqx_stomp_frame:serialize(Frame),
|
|
?LOG(debug, "SEND ~p", [Data]),
|
|
erlang:apply(Fun, [Data] ++ Args),
|
|
{ok, State}.
|
|
|
|
shutdown(_Reason, _State) ->
|
|
ok.
|
|
|
|
timeout(_TRef, {incoming, NewVal},
|
|
State = #pstate{heart_beats = HrtBt}) ->
|
|
case emqx_stomp_heartbeat:check(incoming, NewVal, HrtBt) of
|
|
{error, timeout} ->
|
|
{shutdown, heartbeat_timeout, State};
|
|
{ok, NHrtBt} ->
|
|
{ok, reset_timer(incoming_timer, State#pstate{heart_beats = NHrtBt})}
|
|
end;
|
|
|
|
timeout(_TRef, {outgoing, NewVal},
|
|
State = #pstate{heart_beats = HrtBt,
|
|
heartfun = {Fun, Args}}) ->
|
|
case emqx_stomp_heartbeat:check(outgoing, NewVal, HrtBt) of
|
|
{error, timeout} ->
|
|
_ = erlang:apply(Fun, Args),
|
|
{ok, State};
|
|
{ok, NHrtBt} ->
|
|
{ok, reset_timer(outgoing_timer, State#pstate{heart_beats = NHrtBt})}
|
|
end;
|
|
|
|
timeout(_TRef, clean_trans, State = #pstate{transaction = Trans}) ->
|
|
Now = erlang:system_time(millisecond),
|
|
NTrans = maps:filter(fun(_, {Ts, _}) -> Ts + ?TRANS_TIMEOUT < Now end, Trans),
|
|
{ok, ensure_clean_trans_timer(State#pstate{transaction = NTrans})}.
|
|
|
|
negotiate_version(undefined) ->
|
|
{ok, <<"1.0">>};
|
|
negotiate_version(Accepts) ->
|
|
negotiate_version(?STOMP_VER,
|
|
lists:reverse(
|
|
lists:sort(
|
|
binary:split(Accepts, <<",">>, [global])))).
|
|
|
|
negotiate_version(Ver, []) ->
|
|
{error, <<"Supported protocol versions < ", Ver/binary>>};
|
|
negotiate_version(Ver, [AcceptVer|_]) when Ver >= AcceptVer ->
|
|
{ok, AcceptVer};
|
|
negotiate_version(Ver, [_|T]) ->
|
|
negotiate_version(Ver, T).
|
|
|
|
check_login(undefined, _, AllowAnonymous, _) ->
|
|
AllowAnonymous;
|
|
check_login(_, _, _, undefined) ->
|
|
false;
|
|
check_login(Login, Passcode, _, DefaultUser) ->
|
|
case {list_to_binary(get_value(login, DefaultUser)),
|
|
list_to_binary(get_value(passcode, DefaultUser))} of
|
|
{Login, Passcode} -> true;
|
|
{_, _ } -> false
|
|
end.
|
|
|
|
add_action(Id, Action, ReceiptId, State = #pstate{transaction = Trans}) ->
|
|
case maps:get(Id, Trans, undefined) of
|
|
{Ts, Actions} ->
|
|
NTrans = Trans#{Id => {Ts, [Action|Actions]}},
|
|
{ok, State#pstate{transaction = NTrans}};
|
|
_ ->
|
|
send(error_frame(ReceiptId, ["Transaction ", Id, " not found"]), State)
|
|
end.
|
|
|
|
maybe_send_receipt(undefined, State) ->
|
|
{ok, State};
|
|
maybe_send_receipt(ReceiptId, State) ->
|
|
send(receipt_frame(ReceiptId), State).
|
|
|
|
ack(_Id, State) ->
|
|
State.
|
|
|
|
nack(_Id, State) -> State.
|
|
|
|
header(Name, Headers) ->
|
|
get_value(Name, Headers).
|
|
header(Name, Headers, Val) ->
|
|
get_value(Name, Headers, Val).
|
|
|
|
connected_frame(Headers) ->
|
|
emqx_stomp_frame:make(<<"CONNECTED">>, Headers).
|
|
|
|
receipt_frame(ReceiptId) ->
|
|
emqx_stomp_frame:make(<<"RECEIPT">>, [{<<"receipt-id">>, ReceiptId}]).
|
|
|
|
error_frame(ReceiptId, Msg) ->
|
|
error_frame([{<<"content-type">>, <<"text/plain">>}], ReceiptId, Msg).
|
|
|
|
error_frame(Headers, undefined, Msg) ->
|
|
emqx_stomp_frame:make(<<"ERROR">>, Headers, Msg);
|
|
error_frame(Headers, ReceiptId, Msg) ->
|
|
emqx_stomp_frame:make(<<"ERROR">>, [{<<"receipt-id">>, ReceiptId} | Headers], Msg).
|
|
|
|
next_msgid() ->
|
|
MsgId = case get(msgid) of
|
|
undefined -> 1;
|
|
I -> I
|
|
end,
|
|
put(msgid, MsgId + 1),
|
|
MsgId.
|
|
|
|
next_ackid() ->
|
|
AckId = case get(ackid) of
|
|
undefined -> 1;
|
|
I -> I
|
|
end,
|
|
put(ackid, AckId + 1),
|
|
AckId.
|
|
|
|
make_mqtt_message(Topic, Headers, Body) ->
|
|
Msg = emqx_message:make(stomp, Topic, Body),
|
|
Headers1 = lists:foldl(fun(Key, Headers0) ->
|
|
proplists:delete(Key, Headers0)
|
|
end, Headers, [<<"destination">>,
|
|
<<"content-length">>,
|
|
<<"content-type">>,
|
|
<<"transaction">>,
|
|
<<"receipt">>]),
|
|
emqx_message:set_headers(#{stomp_headers => Headers1}, Msg).
|
|
|
|
receipt_id(Headers) ->
|
|
header(<<"receipt">>, Headers).
|
|
|
|
%%--------------------------------------------------------------------
|
|
%% Transaction Handle
|
|
|
|
handle_recv_send_frame(#stomp_frame{command = <<"SEND">>, headers = Headers, body = Body}, State) ->
|
|
Topic = header(<<"destination">>, Headers),
|
|
_ = maybe_send_receipt(receipt_id(Headers), State),
|
|
_ = emqx_broker:publish(
|
|
make_mqtt_message(Topic, Headers, iolist_to_binary(Body))
|
|
),
|
|
State.
|
|
|
|
handle_recv_ack_frame(#stomp_frame{command = <<"ACK">>, headers = Headers}, State) ->
|
|
Id = header(<<"id">>, Headers),
|
|
_ = maybe_send_receipt(receipt_id(Headers), State),
|
|
ack(Id, State).
|
|
|
|
handle_recv_nack_frame(#stomp_frame{command = <<"NACK">>, headers = Headers}, State) ->
|
|
Id = header(<<"id">>, Headers),
|
|
_ = maybe_send_receipt(receipt_id(Headers), State),
|
|
nack(Id, State).
|
|
|
|
ensure_clean_trans_timer(State = #pstate{transaction = Trans}) ->
|
|
case maps:size(Trans) of
|
|
0 -> State;
|
|
_ -> ensure_timer(clean_trans_timer, State)
|
|
end.
|
|
|
|
%%--------------------------------------------------------------------
|
|
%% Heartbeat
|
|
|
|
parse_heartbeats(Heartbeats) ->
|
|
CxCy = re:split(Heartbeats, <<",">>, [{return, list}]),
|
|
list_to_tuple([list_to_integer(S) || S <- CxCy]).
|
|
|
|
reverse_heartbeats({Cx, Cy}) ->
|
|
iolist_to_binary(io_lib:format("~w,~w", [Cy, Cx])).
|
|
|
|
start_heartbeart_timer(Heartbeats, State) ->
|
|
ensure_timer(
|
|
[incoming_timer, outgoing_timer],
|
|
State#pstate{heart_beats = emqx_stomp_heartbeat:init(Heartbeats)}).
|
|
|
|
%%--------------------------------------------------------------------
|
|
%% Timer
|
|
|
|
ensure_timer([Name], State) ->
|
|
ensure_timer(Name, State);
|
|
ensure_timer([Name | Rest], State) ->
|
|
ensure_timer(Rest, ensure_timer(Name, State));
|
|
|
|
ensure_timer(Name, State = #pstate{timers = Timers}) ->
|
|
TRef = maps:get(Name, Timers, undefined),
|
|
Time = interval(Name, State),
|
|
case TRef == undefined andalso is_integer(Time) andalso Time > 0 of
|
|
true -> ensure_timer(Name, Time, State);
|
|
false -> State %% Timer disabled or exists
|
|
end.
|
|
|
|
ensure_timer(Name, Time, State = #pstate{timers = Timers}) ->
|
|
Msg = maps:get(Name, ?TIMER_TABLE),
|
|
TRef = emqx_misc:start_timer(Time, Msg),
|
|
State#pstate{timers = Timers#{Name => TRef}}.
|
|
|
|
reset_timer(Name, State) ->
|
|
ensure_timer(Name, clean_timer(Name, State)).
|
|
|
|
clean_timer(Name, State = #pstate{timers = Timers}) ->
|
|
State#pstate{timers = maps:remove(Name, Timers)}.
|
|
|
|
interval(incoming_timer, #pstate{heart_beats = HrtBt}) ->
|
|
emqx_stomp_heartbeat:interval(incoming, HrtBt);
|
|
interval(outgoing_timer, #pstate{heart_beats = HrtBt}) ->
|
|
emqx_stomp_heartbeat:interval(outgoing, HrtBt);
|
|
interval(clean_trans_timer, _) ->
|
|
?TRANS_TIMEOUT.
|