1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
|
-module(stomp_worker).
-behaviour(gen_server).
%% API.
-export([start_link/0]).
%% gen_server.
-export([init/1]).
-export([handle_call/3]).
-export([handle_cast/2]).
-export([handle_info/2]).
-export([terminate/2]).
-export([code_change/3]).
%% Testing
-compile(export_all).
-include("stomp.hrl").
%% State of a stomp_worker
-record(state, {connection :: port(),
next_sub :: sub_id(),
subscriptions :: #{ destination() => sub_id() },
subscribers :: #{ destination() => pid() }
}).
-type state() :: #state{}.
%% API implementation
-spec start_link() -> {ok, pid()}.
start_link() ->
{ok, Pid} = gen_server:start_link(?MODULE, [], []),
register(?MODULE, Pid),
{ok, Pid}.
%% gen_server implementation
-spec init(any()) -> {ok, state()}.
init(_Args) ->
%% Fetch configuration from app config
{ok, Host} = application:get_env(stomp, host),
Port = application:get_env(stomp, port, 61613),
Login = application:get_env(stomp, login),
Pass = application:get_env(stomp, passcode),
%% Catch exit signals from linked processes (subscribers dying)
process_flag(trap_exit, true),
%% Establish connection
{ok, Conn} = connect(Host, Port, Login, Pass),
{ok, #state{connection = Conn,
next_sub = 0,
subscriptions = #{},
subscribers = #{}}}.
%% Handle subscription calls
handle_call({subscribe, Dest, Ack}, From, State) ->
%% Subscribe to new destination
SubId = State#state.next_sub,
ok = subscribe(State#state.connection, SubId, Dest, Ack),
%% Add subscription and subscriber to state
Subscriptions = maps:put(SubId, Dest, State#state.subscriptions),
Subscribers = maps:put(SubId, From, State#state.subscribers),
NextSub = SubId + 1,
NewState = State#state{subscriptions = Subscriptions,
subscribers = Subscribers,
next_sub = NextSub },
{reply, {ok, SubId}, NewState};
handle_call(_Req, _From, State) ->
{reply, ignored, State}.
handle_info({tcp, Conn, Frame}, State) when Conn =:= State#state.connection ->
handle_frame(Frame, State);
handle_info(_Msg, State) ->
{noreply, State}.
%% Unused gen_server callbacks
handle_cast(_Msg, State) ->
{noreply, State}.
terminate(_Reason, _State) ->
ok.
code_change(_OldVsn, State, _Extra) ->
{ok, State}.
%% Private functions
-spec connect(list(), integer(), any(), any()) -> {ok, port()}.
connect(Host, Port, Login, Pass) ->
%% STOMP CONNECT frame
Connect = connect_frame(Host, Login, Pass),
%% TODO: Configurable buffer size
%% Frames larger than the user-level buffer will be truncated, so it should
%% never be smaller than the largest expected messages.
{ok, Socket} = gen_tcp:connect(Host, Port, [binary,
{packet, line},
{line_delimiter, $\0},
{buffer, 262144}]),
ok = gen_tcp:send(Socket, Connect),
{ok, Socket}.
-spec subscribe(port(), sub_id(), destination(), ack_mode()) -> ok.
subscribe(Socket, Id, Queue, Ack) ->
{ok, SubscribeFrame} = subscribe_frame(Id, Queue, Ack),
gen_tcp:send(Socket, SubscribeFrame).
%%% Parsing STOMP frames
handle_frame(<<"MESSAGE", "\n", _Frame/binary>>, State) ->
{noreply, State};
handle_frame(Frame, State) ->
io:format("Received unknown frame ~p", [Frame]),
{noreply, State}.
%% Parse out headers into a map
-spec parse_headers(binary()) -> headers().
parse_headers(HeadersBin) ->
Headers = binary:split(HeadersBin, <<"\n">>, [global]),
ToPairs = fun(H, M) -> [K,V | []] = binary:split(H, <<":">>),
maps:put(K, V, M)
end,
{ok, lists:mapfoldl(ToPairs, #{}, Headers)}.
%%% Making STOMP protocol frames
%% Format a header
-spec format_header({binary(), binary()}) -> binary().
format_header({Key, Val}) ->
<<Key/binary, ":", Val/binary, "\n">>.
%% Build a single STOMP frame
-spec make_frame(binary(),
headers(),
binary())
-> {ok, iolist()}.
make_frame(Command, HeaderMap, Body) ->
Headers = lists:map(fun format_header/1, maps:to_list(HeaderMap)),
Frame = [Command, <<"\n">>, Headers, <<"\n">>, Body, <<0>>],
{ok, Frame}.
%%% Default frames
-spec connect_frame(list(), any(), any()) -> iolist().
connect_frame(Host, {ok, Login}, {ok, Pass}) ->
make_frame(<<"CONNECT">>,
#{<<"accept-version">> => <<"1.2">>,
<<"host">> => Host,
<<"login">> => Login,
<<"passcode">> => Pass,
<<"heart-beat">> => <<"0,5000">>},
[]);
connect_frame(Host, _Login, _Pass) ->
make_frame(<<"CONNECT">>,
#{<<"accept-version">> => <<"1.2">>,
<<"host">> => Host,
%% Expect a server heartbeat every 5 seconds, let the server
%% expect one every 10. We don't actually check this and just
%% echo server heartbeats.
%% TODO: For now the server is told not to expect replies due to
%% a weird behaviour.
<<"heart-beat">> => <<"0,5000">>},
[]).
-spec subscribe_frame(sub_id(), destination(), ack_mode()) -> iolist().
subscribe_frame(Id, Queue, Ack) ->
make_frame(<<"SUBSCRIBE">>,
#{<<"id">> => integer_to_binary(Id),
<<"destination">> => Queue,
<<"ack">> => ack_mode_to_binary(Ack)},
[]).
-spec ack_mode_to_binary(ack_mode()) -> binary().
ack_mode_to_binary(AckMode) ->
case AckMode of
auto -> <<"auto">>;
client -> <<"client">>;
client_individual -> <<"client-individual">>
end.
%% -spec ack_frame(binary()) -> iolist().
%% ack_frame(MessageID) ->
%% make_frame(<<"ACK">>,
%% [{"id", MessageID}],
%% []).
|
