emqx / esockd / 433

Committed 13 Jul 2024 09:38AM UTC coverage: 67.152%. First build

Build # 433

Build Type

Pull #190

github

Committed by

web-flow

Commit Message

Merge 2da91fb66 into 4ec038215

Pull Request Pull Request #190: fix: fixed the proxy could not be released even if the RC was 0

Run Details

0 of 43 new or added lines in 3 files covered. (0.0%)

877 of 1306 relevant lines covered (67.15%)

56.79 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

/src/udp_proxy/esockd_udp_proxy.erl

%%--------------------------------------------------------------------
%% Copyright (c) 2020 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.
%%--------------------------------------------------------------------

-module(esockd_udp_proxy).

-behaviour(gen_server).

-include("include/esockd_proxy.hrl").

%% API
-export([start_link/3, send/2, close/1, takeover/2]).

%% gen_server callbacks
-export([
    init/1,
    handle_call/3,
    handle_cast/2,
    handle_info/2,
    terminate/2
]).

-export_type([connection_options/0]).

-define(NOW, erlang:system_time(second)).
-define(ERROR_MSG(Format, Args),
    error_logger:error_msg("[~s]: " ++ Format, [?MODULE | Args])
).
-define(DEF_HEARTBEAT, 60).

-type timespan() :: non_neg_integer().

%%--------------------------------------------------------------------
%%  Definitions
%%--------------------------------------------------------------------

-type state() :: #{
    connection_mod := connection_module(),
    connection_id := connection_id() | undefined,
    connection_state := connection_state(),
    connection_pid := pid() | undefined,
    connection_ref := reference() | undefined,
    connection_options := connection_options(),
    %% last source's connection active time
    last_time := pos_integer(),
    transport := proxy_transport(),
    peer := peer()
}.

%%--------------------------------------------------------------------
%%- API
%%--------------------------------------------------------------------

start_link(Transport, Peer, Opts) ->
    gen_server:start_link(?MODULE, [Transport, Peer, Opts], []).

-spec send(proxy_id(), binary()) -> ok.
send(ProxyId, Data) ->
    gen_server:cast(ProxyId, {send, Data}).

close(ProxyId) ->
    case erlang:is_process_alive(ProxyId) of
        true ->
            gen_server:call(ProxyId, close);
        _ ->
            ok
    end.

takeover(ProxyId, CId) ->
    _ = gen_server:cast(ProxyId, {?FUNCTION_NAME, CId}),
    ok.

%%--------------------------------------------------------------------
%%- gen_server callbacks
%%--------------------------------------------------------------------

init([Transport, Peer, #{esockd_proxy_opts := Opts} = COpts]) ->
    #{connection_mod := Mod} = Opts,
    heartbeat(maps:get(heartbeat, Opts, ?DEF_HEARTBEAT)),
    init_transport(Transport, Peer, #{
        last_time => ?NOW,
        connection_mod => Mod,
        connection_options => COpts,
        connection_state => esockd_udp_proxy_connection:initialize(Mod, COpts),
        connection_id => undefined,
        connection_pid => undefined,
        connection_ref => undefined
    }).

handle_call(close, _From, State) ->
    {stop, {shutdown, close_transport}, ok, State};
handle_call(Request, _From, State) ->
    ?ERROR_MSG("Unexpected call: ~p", [Request]),
    {reply, ok, State}.

handle_cast({send, Data}, #{transport := Transport, peer := Peer} = State) ->
    case send(Transport, Peer, Data) of
        ok ->
            {noreply, State};
        {error, Reason} ->
            ?ERROR_MSG("Send failed, Reason: ~0p", [Reason]),
            {stop, {sock_error, Reason}, State}
    end;
handle_cast({takeover, CId}, #{connection_id := CId} = State) ->
    {stop, {shutdown, takeover}, State};
handle_cast({takeover, _CId}, State) ->
    {noreply, State};
handle_cast(Request, State) ->
    ?ERROR_MSG("Unexpected cast: ~p", [Request]),
    {noreply, State}.

handle_info({datagram, _SockPid, Data}, State) ->
    handle_incoming(Data, State);
handle_info({ssl, _Socket, Data}, State) ->
    handle_incoming(Data, State);
handle_info({heartbeat, Span}, #{last_time := LastTime} = State) ->
    Now = ?NOW,
    case Now - LastTime > Span of
        true ->
            {stop, normal, State};
        _ ->
            heartbeat(Span),
            {noreply, State}
    end;
handle_info({ssl_error, _Sock, Reason}, State) ->
    {stop, Reason, socket_exit(State)};
handle_info({ssl_closed, _Sock}, State) ->
    {stop, ssl_closed, socket_exit(State)};
handle_info(
    {'DOWN', _, process, Pid, _Reason},
    State
) ->
    ct:print(" >>>> conn down, Pid:~p, Reason:~p~n", [Pid, _Reason]),
    {stop, {shutdown, connection_closed}, State};
handle_info(Info, State) ->
    ?ERROR_MSG("Unexpected info: ~p", [Info]),
    {noreply, State}.

terminate(Reason, #{transport := Transport} = State) ->
    close_transport(Transport),
    Clear =
        case Reason of
            close_transport ->
                false;
            connection_closed ->
                false;
            takeover ->
                false;
            _ ->
                true
        end,
    ct:print(">>>> esockd:~p terminate:~p~n", [self(), Reason]),
    detach(State, Clear).

%%--------------------------------------------------------------------
%%- Internal functions
%%--------------------------------------------------------------------
-spec handle_incoming(socket_packet(), state()) -> state().
handle_incoming(
    Data,
    #{transport := Transport, peer := Peer, connection_mod := Mod, connection_state := CState} =
        State
) ->
    State2 = State#{last_time := ?NOW},
    case esockd_udp_proxy_connection:get_connection_id(Mod, Transport, Peer, CState, Data) of
        {ok, CId, Packet, CState2} ->
            dispatch(Mod, CId, Data, Packet, State2#{connection_state := CState2});
        {error, Reply} ->
            ?ERROR_MSG("Can't get connection id, Transport:~0p, Peer:~0p, Mod:~0p", [
                Transport, Peer, Mod
            ]),
            send(Transport, Peer, Reply),
            {stop, {shutdown, no_clientid}, State2};
        invalid ->
            ?ERROR_MSG("Can't get connection id, Transport:~0p, Peer:~0p, Mod:~0p", [
                Transport, Peer, Mod
            ]),
            {stop, {shutdown, no_clientid}, State2}
    end.

-spec dispatch(
    connection_module(),
    esockd_transport:socket(),
    connection_id(),
    connection_packet(),
    state()
) ->
    state().
dispatch(
    Mod,
    CId,
    Data,
    Packet,
    #{
        transport := Transport,
        connection_state := CState
    } =
        State
) ->
    case lookup(CId, State) of
        {ok, Pid} ->
            Result = attach(CId, State, Pid),
            esockd_udp_proxy_connection:dispatch(
                Mod, Pid, CState, {Transport, Data, Packet}
            ),
            {noreply, Result};
        {error, Reason} ->
            ?ERROR_MSG("Dispatch failed, Reason:~0p", [Reason]),
            {noreply, State}
    end.

-spec attach(connection_id(), state(), pid()) -> state().
attach(CId, #{connection_mod := Mod, connection_id := undefined} = State, Pid) ->
    esockd_udp_proxy_db:attach(Mod, CId),
    Ref = erlang:monitor(process, Pid),
    State#{connection_id := CId, connection_pid := Pid, connection_ref := Ref};
attach(CId, #{connection_id := OldId} = State, Pid) when CId =/= OldId ->
    State2 = detach(State, false),
    attach(CId, State2, Pid);
attach(_CId, State, _Pid) ->
    State.

detach(State) ->
    detach(State, true).

-spec detach(state()) -> state().
detach(#{connection_id := undefined} = State, _Clear) ->
    State;
detach(
    #{
        connection_id := CId,
        connection_pid := Pid,
        connection_ref := Ref,
        connection_mod := Mod,
        connection_state := CState
    } = State,
    Clear
) ->
    erlang:demonitor(Ref),

    Result = esockd_udp_proxy_db:detach(Mod, CId),
    ct:print(">>>> detach:~p result:~p, pid:~p~n", [self(), Result, Pid]),
    case Clear andalso Result of
        true ->
            case erlang:is_process_alive(Pid) of
                true ->
                    esockd_udp_proxy_connection:close(Mod, Pid, CState);
                _ ->
                    ok
            end;
        _ ->
            ok
    end,
    State#{connection_id := undefined, connection_pid := undefined, connection_ref := undefined}.

-spec socket_exit(state()) -> state().
socket_exit(State) ->
    detach(State).

-spec heartbeat(timespan()) -> ok.
heartbeat(Span) ->
    erlang:send_after(timer:seconds(Span), self(), {?FUNCTION_NAME, Span}),
    ok.

-spec lookup(connection_id(), state()) -> {ok, pid()} | {error, Reason :: term()}.
lookup(_CId, #{connection_pid := Pid}) when is_pid(Pid) ->
    {ok, Pid};
lookup(CId, #{
    connection_pid := undefined,
    connection_mod := Mod,
    transport := Transport,
    peer := Peer,
    connection_options := Opts
}) ->
    case esockd_udp_proxy_connection:find_or_create(Mod, CId, Transport, Peer, Opts) of
        {ok, Pid} ->
            {ok, Pid};
        ignore ->
            {error, ignore};
        Error ->
            Error
    end.

-spec send(proxy_transport(), peer(), binary()) -> _.
send({?PROXY_TRANSPORT, _, Socket}, {IP, Port}, Data) when is_port(Socket) ->
    gen_udp:send(Socket, IP, Port, Data);
send({?PROXY_TRANSPORT, _, Socket}, _Peer, Data) ->
    esockd_transport:send(Socket, Data).

init_transport({udp, _, Sock}, Peer, State) ->
    {ok, State#{
        transport => {?PROXY_TRANSPORT, self(), Sock},
        peer => Peer
    }};
init_transport(esockd_transport, Sock, State) ->
    case esockd_transport:wait(Sock) of
        {ok, NSock} ->
            {ok, State#{
                transport => {?PROXY_TRANSPORT, self(), NSock},
                peer => esockd_transport:peername(NSock)
            }};
        Error ->
            Error
    end.

close_transport({?PROXY_TRANSPORT, _, Sock}) when is_port(Sock) ->
    ok;
close_transport({?PROXY_TRANSPORT, _, Sock}) ->
    esockd_transport:fast_close(Sock).

1	%%--------------------------------------------------------------------
2	%% Copyright (c) 2020 EMQ Technologies Co., Ltd. All Rights Reserved.
3	%%
4	%% Licensed under the Apache License, Version 2.0 (the "License");
5	%% you may not use this file except in compliance with the License.
6	%% You may obtain a copy of the License at
7	%%
8	%% http://www.apache.org/licenses/LICENSE-2.0
9	%%
10	%% Unless required by applicable law or agreed to in writing, software
11	%% distributed under the License is distributed on an "AS IS" BASIS,
12	%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	%% See the License for the specific language governing permissions and
14	%% limitations under the License.
15	%%--------------------------------------------------------------------
16
17	-module(esockd_udp_proxy).
18
19	-behaviour(gen_server).
20
21	-include("include/esockd_proxy.hrl").
22
23	%% API
24	-export([start_link/3, send/2, close/1, takeover/2]).
25
26	%% gen_server callbacks
27	-export([
28	init/1,
29	handle_call/3,
30	handle_cast/2,
31	handle_info/2,
32	terminate/2
33	]).
34
35	-export_type([connection_options/0]).
36
37	-define(NOW, erlang:system_time(second)).
38	-define(ERROR_MSG(Format, Args),
39	error_logger:error_msg("[~s]: " ++ Format, [?MODULE \| Args])
40	).
41	-define(DEF_HEARTBEAT, 60).
42
43	-type timespan() :: non_neg_integer().
44
45	%%--------------------------------------------------------------------
46	%% Definitions
47	%%--------------------------------------------------------------------
48
49	-type state() :: #{
50	connection_mod := connection_module(),
51	connection_id := connection_id() \| undefined,
52	connection_state := connection_state(),
53	connection_pid := pid() \| undefined,
54	connection_ref := reference() \| undefined,
55	connection_options := connection_options(),
56	%% last source's connection active time
57	last_time := pos_integer(),
58	transport := proxy_transport(),
59	peer := peer()
60	}.
61
62	%%--------------------------------------------------------------------
63	%%- API
64	%%--------------------------------------------------------------------
65
66	start_link(Transport, Peer, Opts) ->
67	gen_server:start_link(?MODULE, [Transport, Peer, Opts], []).	×
68
69	-spec send(proxy_id(), binary()) -> ok.
70	send(ProxyId, Data) ->
71	gen_server:cast(ProxyId, {send, Data}).	×
72
73	close(ProxyId) ->
74	case erlang:is_process_alive(ProxyId) of	×
75	true ->
76	gen_server:call(ProxyId, close);	×
77	_ ->
78	ok	×
79	end.
80
81	takeover(ProxyId, CId) ->
NEW 82	_ = gen_server:cast(ProxyId, {?FUNCTION_NAME, CId}),	×
NEW 83	ok.	×
84
85	%%--------------------------------------------------------------------
86	%%- gen_server callbacks
87	%%--------------------------------------------------------------------
88
89	init([Transport, Peer, #{esockd_proxy_opts := Opts} = COpts]) ->
90	#{connection_mod := Mod} = Opts,	×
91	heartbeat(maps:get(heartbeat, Opts, ?DEF_HEARTBEAT)),	×
92	init_transport(Transport, Peer, #{	×
93	last_time => ?NOW,
94	connection_mod => Mod,
95	connection_options => COpts,
96	connection_state => esockd_udp_proxy_connection:initialize(Mod, COpts),
97	connection_id => undefined,
98	connection_pid => undefined,
99	connection_ref => undefined
100	}).
101
102	handle_call(close, _From, State) ->
103	{stop, {shutdown, close_transport}, ok, State};	×
104	handle_call(Request, _From, State) ->
105	?ERROR_MSG("Unexpected call: ~p", [Request]),	×
106	{reply, ok, State}.	×
107
108	handle_cast({send, Data}, #{transport := Transport, peer := Peer} = State) ->
109	case send(Transport, Peer, Data) of	×
110	ok ->
111	{noreply, State};	×
112	{error, Reason} ->
113	?ERROR_MSG("Send failed, Reason: ~0p", [Reason]),	×
114	{stop, {sock_error, Reason}, State}	×
115	end;
116	handle_cast({takeover, CId}, #{connection_id := CId} = State) ->
NEW 117	{stop, {shutdown, takeover}, State};	×
118	handle_cast({takeover, _CId}, State) ->
NEW 119	{noreply, State};	×
120	handle_cast(Request, State) ->
121	?ERROR_MSG("Unexpected cast: ~p", [Request]),	×
122	{noreply, State}.	×
123
124	handle_info({datagram, _SockPid, Data}, State) ->
NEW 125	handle_incoming(Data, State);	×
126	handle_info({ssl, _Socket, Data}, State) ->
NEW 127	handle_incoming(Data, State);	×
128	handle_info({heartbeat, Span}, #{last_time := LastTime} = State) ->
129	Now = ?NOW,	×
130	case Now - LastTime > Span of	×
131	true ->
132	{stop, normal, State};	×
133	_ ->
134	heartbeat(Span),	×
135	{noreply, State}	×
136	end;
137	handle_info({ssl_error, _Sock, Reason}, State) ->
138	{stop, Reason, socket_exit(State)};	×
139	handle_info({ssl_closed, _Sock}, State) ->
140	{stop, ssl_closed, socket_exit(State)};	×
141	handle_info(
142	{'DOWN', _, process, Pid, _Reason},
143	State
144	) ->
NEW 145	ct:print(" >>>> conn down, Pid:~p, Reason:~p~n", [Pid, _Reason]),	×
146	{stop, {shutdown, connection_closed}, State};	×
147	handle_info(Info, State) ->
148	?ERROR_MSG("Unexpected info: ~p", [Info]),	×
149	{noreply, State}.	×
150
151	terminate(Reason, #{transport := Transport} = State) ->
152	close_transport(Transport),	×
153	Clear =	×
154	case Reason of
155	close_transport ->
156	false;	×
157	connection_closed ->
158	false;	×
159	takeover ->
NEW 160	false;	×
161	_ ->
162	true	×
163	end,
NEW 164	ct:print(">>>> esockd:~p terminate:~p~n", [self(), Reason]),	×
165	detach(State, Clear).	×
166
167	%%--------------------------------------------------------------------
168	%%- Internal functions
169	%%--------------------------------------------------------------------
170	-spec handle_incoming(socket_packet(), state()) -> state().
171	handle_incoming(
172	Data,
173	#{transport := Transport, peer := Peer, connection_mod := Mod, connection_state := CState} =
174	State
175	) ->
176	State2 = State#{last_time := ?NOW},	×
177	case esockd_udp_proxy_connection:get_connection_id(Mod, Transport, Peer, CState, Data) of	×
178	{ok, CId, Packet, CState2} ->
179	dispatch(Mod, CId, Data, Packet, State2#{connection_state := CState2});	×
180	{error, Reply} ->
NEW 181	?ERROR_MSG("Can't get connection id, Transport:~0p, Peer:~0p, Mod:~0p", [	×
182	Transport, Peer, Mod
183	]),
NEW 184	send(Transport, Peer, Reply),	×
NEW 185	{stop, {shutdown, no_clientid}, State2};	×
186	invalid ->
187	?ERROR_MSG("Can't get connection id, Transport:~0p, Peer:~0p, Mod:~0p", [	×
188	Transport, Peer, Mod
189	]),
NEW 190	{stop, {shutdown, no_clientid}, State2}	×
191	end.
192
193	-spec dispatch(
194	connection_module(),
195	esockd_transport:socket(),
196	connection_id(),
197	connection_packet(),
198	state()
199	) ->
200	state().
201	dispatch(
202	Mod,
203	CId,
204	Data,
205	Packet,
206	#{
207	transport := Transport,
208	connection_state := CState
209	} =
210	State
211	) ->
NEW 212	case lookup(CId, State) of	×
213	{ok, Pid} ->
NEW 214	Result = attach(CId, State, Pid),	×
215	esockd_udp_proxy_connection:dispatch(	×
216	Mod, Pid, CState, {Transport, Data, Packet}
217	),
NEW 218	{noreply, Result};	×
219	{error, Reason} ->
220	?ERROR_MSG("Dispatch failed, Reason:~0p", [Reason]),	×
NEW 221	{noreply, State}	×
222	end.
223
224	-spec attach(connection_id(), state(), pid()) -> state().
225	attach(CId, #{connection_mod := Mod, connection_id := undefined} = State, Pid) ->
226	esockd_udp_proxy_db:attach(Mod, CId),	×
NEW 227	Ref = erlang:monitor(process, Pid),	×
NEW 228	State#{connection_id := CId, connection_pid := Pid, connection_ref := Ref};	×
229	attach(CId, #{connection_id := OldId} = State, Pid) when CId =/= OldId ->
NEW 230	State2 = detach(State, false),	×
NEW 231	attach(CId, State2, Pid);	×
232	attach(_CId, State, _Pid) ->
233	State.	×
234
235	detach(State) ->
236	detach(State, true).	×
237
238	-spec detach(state()) -> state().
239	detach(#{connection_id := undefined} = State, _Clear) ->
240	State;	×
241	detach(
242	#{
243	connection_id := CId,
244	connection_pid := Pid,
245	connection_ref := Ref,
246	connection_mod := Mod,
247	connection_state := CState
248	} = State,
249	Clear
250	) ->
NEW 251	erlang:demonitor(Ref),	×
252
NEW 253	Result = esockd_udp_proxy_db:detach(Mod, CId),	×
NEW 254	ct:print(">>>> detach:~p result:~p, pid:~p~n", [self(), Result, Pid]),	×
NEW 255	case Clear andalso Result of	×
256	true ->
257	case erlang:is_process_alive(Pid) of	×
258	true ->
259	esockd_udp_proxy_connection:close(Mod, Pid, CState);	×
260	_ ->
261	ok	×
262	end;
263	_ ->
264	ok	×
265	end,
NEW 266	State#{connection_id := undefined, connection_pid := undefined, connection_ref := undefined}.	×
267
268	-spec socket_exit(state()) -> state().
269	socket_exit(State) ->
270	detach(State).	×
271
272	-spec heartbeat(timespan()) -> ok.
273	heartbeat(Span) ->
274	erlang:send_after(timer:seconds(Span), self(), {?FUNCTION_NAME, Span}),	×
275	ok.	×
276
277	-spec lookup(connection_id(), state()) -> {ok, pid()} \| {error, Reason :: term()}.
278	lookup(_CId, #{connection_pid := Pid}) when is_pid(Pid) ->
NEW 279	{ok, Pid};	×
280	lookup(CId, #{
281	connection_pid := undefined,
282	connection_mod := Mod,
283	transport := Transport,
284	peer := Peer,
285	connection_options := Opts
286	}) ->
NEW 287	case esockd_udp_proxy_connection:find_or_create(Mod, CId, Transport, Peer, Opts) of	×
288	{ok, Pid} ->
NEW 289	{ok, Pid};	×
290	ignore ->
NEW 291	{error, ignore};	×
292	Error ->
NEW 293	Error	×
294	end.
295
296	-spec send(proxy_transport(), peer(), binary()) -> _.
297	send({?PROXY_TRANSPORT, _, Socket}, {IP, Port}, Data) when is_port(Socket) ->
298	gen_udp:send(Socket, IP, Port, Data);	×
299	send({?PROXY_TRANSPORT, _, Socket}, _Peer, Data) ->
300	esockd_transport:send(Socket, Data).	×
301
302	init_transport({udp, _, Sock}, Peer, State) ->
303	{ok, State#{	×
304	transport => {?PROXY_TRANSPORT, self(), Sock},
305	peer => Peer
306	}};
307	init_transport(esockd_transport, Sock, State) ->
308	case esockd_transport:wait(Sock) of	×
309	{ok, NSock} ->
310	{ok, State#{	×
311	transport => {?PROXY_TRANSPORT, self(), NSock},
312	peer => esockd_transport:peername(NSock)
313	}};
314	Error ->
315	Error	×
316	end.
317
318	close_transport({?PROXY_TRANSPORT, _, Sock}) when is_port(Sock) ->
319	ok;	×
320	close_transport({?PROXY_TRANSPORT, _, Sock}) ->
321	esockd_transport:fast_close(Sock).	×

emqx / esockd / 433

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous