emqx / esockd / 349

Committed 28 Nov 2023 12:32PM UTC coverage: 70.65% (-0.1%) from 70.781%

Build # 349

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #182 from emqx/dev/william/improve-robustness

improve robustness

Run Details

9 of 13 new or added lines in 4 files covered. (69.23%)

3 existing lines in 2 files now uncovered.

739 of 1046 relevant lines covered (70.65%)

58.94 hits per line

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

54.04

/src/esockd_transport.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_transport).

-include("esockd.hrl").
-include_lib("ssl/src/ssl_api.hrl").

-export([type/1, is_ssl/1]).
-export([listen/2]).
-export([ready/3, wait/1]).
-export([send/2, async_send/2, async_send/3, recv/2, recv/3, async_recv/2, async_recv/3]).
-export([controlling_process/2]).
-export([close/1, fast_close/1]).
-export([getopts/2, setopts/2, getstat/2]).
-export([sockname/1, peername/1, shutdown/2]).
-export([peercert/1, peer_cert_subject/1, peer_cert_common_name/1, peersni/1]).
-export([ensure_ok_or_exit/2]).
-export([gc/1]).

%% Internal callbacks
-export([ ssl_upgrade_fun/1
        , ssl_upgrade/3
        , proxy_upgrade_fun/1
        , proxy_upgrade/2
        ]).

-export_type([socket/0]).

-type(ssl_socket() :: #ssl_socket{}).
-type(proxy_socket() :: #proxy_socket{}).
-type(socket() :: inet:socket() | ssl_socket() | proxy_socket() | #sslsocket{}).

-spec(type(socket()) -> tcp | ssl | proxy).
type(Sock) when is_port(Sock) ->
    tcp;
type(#ssl_socket{})  ->
    ssl;
type(#proxy_socket{}) ->
    proxy.

-spec(is_ssl(socket()) -> boolean()).
is_ssl(Sock) when is_port(Sock) ->
    false;
is_ssl(#ssl_socket{})  ->
    true;
is_ssl(#proxy_socket{socket = Sock}) ->
    is_ssl(Sock).

-spec(ready(pid(), socket(), [esockd:sock_fun()]) -> any()).
ready(Pid, Sock, UpgradeFuns) ->
    Pid ! {sock_ready, Sock, UpgradeFuns}.

-spec(wait(socket()) -> {ok, socket()} | {error, term()}).
wait(Sock) ->
    receive
        {sock_ready, Sock, UpgradeFuns} ->
            upgrade(Sock, UpgradeFuns)
    end.

-spec(upgrade(socket(), [esockd:sock_fun()]) -> {ok, socket()} | {error, term()}).
upgrade(Sock, []) ->
    {ok, Sock};
upgrade(Sock, [{Fun, Args}|More]) ->
    case apply(Fun, [Sock|Args]) of
        {ok, NewSock} -> upgrade(NewSock, More);
        Error         -> fast_close(Sock), Error
    end.

-spec(listen(inet:port_number(), [gen_tcp:listen_option()])
      -> {ok, inet:socket()} | {error, system_limit | inet:posix()}).
listen(Port, Opts) ->
    gen_tcp:listen(Port, Opts).

-spec(controlling_process(socket(), pid()) -> ok | {error, Reason} when
      Reason :: closed | not_owner | badarg | inet:posix()).
controlling_process(Sock, NewOwner) when is_port(Sock) ->
    gen_tcp:controlling_process(Sock, NewOwner);
controlling_process(#ssl_socket{ssl = SslSock}, NewOwner) ->
    ssl:controlling_process(SslSock, NewOwner);
controlling_process(#proxy_socket{socket = Sock}, NewOwner) ->
    controlling_process(Sock, NewOwner);
%% Before upgrading, the DTLS socket is #sslsocket{}
%% type, instead of a port().
%%
%% See: ssl:transport_accept/1
controlling_process(SslSock = #sslsocket{}, NewOwner) ->
    ssl:controlling_process(SslSock, NewOwner).

-spec(close(socket()) -> ok | {error, term()}).
close(Sock) when is_port(Sock) ->
    gen_tcp:close(Sock);
close(#ssl_socket{ssl = SslSock}) ->
    ssl:close(SslSock);
close(#proxy_socket{socket = Sock}) ->
    close(Sock).

-spec(fast_close(socket()) -> ok).
fast_close(Sock) when is_port(Sock) ->
    catch port_close(Sock), ok;
fast_close(#ssl_socket{tcp = Sock, ssl = SslSock}) ->
    _ = fast_close_sslsock(SslSock),
    catch port_close(Sock), ok;
fast_close(SslSock = #sslsocket{}) ->
    fast_close_sslsock(SslSock);
fast_close(#proxy_socket{socket = Sock}) ->
    fast_close(Sock).

%% @private
fast_close_sslsock(SslSock) ->
    {Pid, MRef} = spawn_monitor(fun() -> ssl:close(SslSock) end),
    TRef = erlang:send_after(?SSL_CLOSE_TIMEOUT, self(), {Pid, ssl_close_timeout}),
    _ = receive
            {Pid, ssl_close_timeout} ->
                erlang:demonitor(MRef, [flush]),
                exit(Pid, kill);
            {'DOWN', MRef, process, Pid, _Reason} ->
                erlang:cancel_timer(TRef)
        end,
    ok.

-spec(send(socket(), iodata()) -> ok | {error, Reason} when
      Reason :: closed | timeout | inet:posix()).
send(Sock, Data) when is_port(Sock) ->
    gen_tcp:send(Sock, Data);
send(#ssl_socket{ssl = SslSock}, Data) ->
    ssl:send(SslSock, Data);
send(#proxy_socket{socket = Sock}, Data) ->
    send(Sock, Data).

%% @doc Port command to write data.
-spec(async_send(socket(), iodata()) -> ok | {error, Reason} when
      Reason :: close | timeout | inet:posix()).
async_send(Sock, Data) ->
    async_send(Sock, Data, []).

-spec(async_send(socket(), iodata(), [force|nosuspend]) -> ok | {error, Reason} when
      Reason :: close | timeout | inet:posix()).
async_send(Sock, Data, Options) when is_port(Sock) ->
    do_port_command(Sock, Data, Options);
async_send(Sock = #ssl_socket{ssl = SslSock}, Data, _Options) ->
    case ssl:send(SslSock, Data) of
        ok ->
            %% TODO: get rid of this message send
            self() ! {inet_reply, Sock, ok},
            ok;
        Error ->
            Error
    end;
async_send(#proxy_socket{socket = Sock}, Data, _Options) ->
    async_send(Sock, Data).

-if(?OTP_RELEASE >= 26).
do_port_command(Sock, Data, _Options) ->
    case gen_tcp:send(Sock, Data) of
        ok ->
            %% TODO: get rid of this message send
            self() ! {inet_reply, Sock, ok},
            ok;
        {error, Reason} ->
            {error, Reason}
    end.
-else.
do_port_command(Sock, Data, Options) ->
    try erlang:port_command(Sock, Data, Options) of
        true ->
            ok;
        false -> %% nosuspend option and port busy
            {error, busy}
    catch
        error:_Error ->
            {error, einval}
    end.
-endif.

-spec(recv(socket(), non_neg_integer())
      -> {ok, iodata()} | {error, closed | inet:posix()}).
recv(Sock, Length) when is_port(Sock) ->
    gen_tcp:recv(Sock, Length);
recv(#ssl_socket{ssl = SslSock}, Length) ->
    ssl:recv(SslSock, Length);
recv(#proxy_socket{socket = Sock}, Length) ->
    recv(Sock, Length).

-spec(recv(socket(), non_neg_integer(), timeout())
      -> {ok, iodata()} | {error, closed | inet:posix()}).
recv(Sock, Length, Timeout) when is_port(Sock) ->
    gen_tcp:recv(Sock, Length, Timeout);
recv(#ssl_socket{ssl = SslSock}, Length, Timeout)  ->
    ssl:recv(SslSock, Length, Timeout);
recv(#proxy_socket{socket = Sock}, Length, Timeout) ->
    recv(Sock, Length, Timeout).

%% @doc Async receive data.
-spec(async_recv(socket(), non_neg_integer()) -> {ok, reference()}).
async_recv(Sock, Length) ->
    async_recv(Sock, Length, infinity).

-spec(async_recv(socket(), non_neg_integer(), timeout()) -> {ok, reference()}).
async_recv(Sock = #ssl_socket{ssl = SslSock}, Length, Timeout) ->
    Self = self(),
    Ref = make_ref(),
    spawn(fun() ->
              Self ! {inet_async, Sock, Ref, ssl:recv(SslSock, Length, Timeout)}
          end),
    {ok, Ref};
async_recv(Sock, Length, infinity) when is_port(Sock) ->
    prim_inet:async_recv(Sock, Length, -1);
async_recv(Sock, Length, Timeout) when is_port(Sock) ->
    prim_inet:async_recv(Sock, Length, Timeout);
async_recv(#proxy_socket{socket = Sock}, Length, Timeout) ->
    async_recv(Sock, Length, Timeout).

%% @doc Get socket options.
-spec(getopts(socket(), [inet:socket_getopt()])
      -> {ok, [inet:socket_setopt()]} | {error, inet:posix()}).
getopts(Sock, OptionNames) when is_port(Sock) ->
    inet:getopts(Sock, OptionNames);
getopts(#ssl_socket{ssl = SslSock}, OptionNames) ->
    ssl:getopts(SslSock, OptionNames);
getopts(SslSock = #sslsocket{}, OptionNames) ->
    ssl:getopts(SslSock, OptionNames);
getopts(#proxy_socket{socket = Sock}, OptionNames) ->
    getopts(Sock, OptionNames).

%% @doc Set socket options
-spec(setopts(socket(), [inet:socket_setopt()]) -> ok | {error, inet:posix()}).
setopts(Sock, Opts) when is_port(Sock) ->
    inet:setopts(Sock, Opts);
setopts(#ssl_socket{ssl = SslSock}, Opts) ->
    ssl:setopts(SslSock, Opts);
setopts(SslSock = #sslsocket{}, Opts) ->
    ssl:setopts(SslSock, Opts);
setopts(#proxy_socket{socket = Socket}, Opts) ->
    setopts(Socket, Opts).

%% @doc Get socket stats
-spec(getstat(socket(), [inet:stat_option()])
      -> {ok, [{inet:stat_option(), integer()}]} | {error, inet:posix()}).
getstat(Sock, Stats) when is_port(Sock) ->
    inet:getstat(Sock, Stats);
getstat(#ssl_socket{tcp = Sock}, Stats) when is_port(Sock) ->
    inet:getstat(Sock, Stats);
getstat(#ssl_socket{ssl = SslSock}, Stats) ->
    ssl_getstat(SslSock, Stats);
getstat(SslSock = #sslsocket{}, Stats) ->
    ssl_getstat(SslSock, Stats);
getstat(#proxy_socket{socket = Sock}, Stats) ->
    getstat(Sock, Stats).

%% Below OTP 23, the ssl:getstats/2 will get a funcation_clause
%% exception for dtls socket
-if(?OTP_RELEASE >= 23).
ssl_getstat(SslSock, Stats) ->
    ssl:getstat(SslSock, Stats).
-else.
ssl_getstat(#sslsocket{fd = {gen_udp, {_,{{_sockhost,_sockport}, SockPort}},_}}, Stats) ->
    inet:getstat(SockPort, Stats);
ssl_getstat(SslSock, Stats) ->
    ssl:getstat(SslSock, Stats).
-endif.

%% @doc Sockname
-spec(sockname(socket()) -> {ok, {inet:ip_address(), inet:port_number()}} |
                          {error, inet:posix()}).
sockname(Sock) when is_port(Sock) ->
    inet:sockname(Sock);
sockname(#ssl_socket{ssl = SslSock}) ->
    ssl:sockname(SslSock);
sockname(SslSock = #sslsocket{}) ->
    ssl:sockname(SslSock);
sockname(#proxy_socket{dst_addr = DstAddr, dst_port = DstPort}) ->
    {ok, {DstAddr, DstPort}}.

%% @doc Peername
-spec(peername(socket()) -> {ok, {inet:ip_address(), inet:port_number()}} |
                          {error, inet:posix()}).
peername(Sock) when is_port(Sock) ->
    inet:peername(Sock);
peername(#ssl_socket{ssl = SslSock}) ->
    ssl:peername(SslSock);
peername(#proxy_socket{src_addr = SrcAddr, src_port = SrcPort}) ->
    {ok, {SrcAddr, SrcPort}};
%% Before upgrading, the DTLS socket is #sslsocket{}
%% type, instead of a port().
%%
%% See: ssl:transport_accept/1
peername(SslSock = #sslsocket{}) ->
    ssl:peername(SslSock).

%% @doc Socket peercert
-spec(peercert(socket()) -> nossl | binary() | list(pp2_additional_ssl_field()) |
                          {error, term()}).
peercert(Sock) when is_port(Sock) ->
    nossl;
peercert(#ssl_socket{ssl = SslSock}) ->
    case ssl:peercert(SslSock) of
        {ok, Cert} -> Cert;
        %% One-way SSL
        {error, no_peercert} ->
            undefined;
        Error -> Error
    end;
peercert(SslSock = #sslsocket{}) ->
    case ssl:peercert(SslSock) of
        {ok, Cert} -> Cert;
        {error, no_peercert} -> undefined;
        Error -> Error
    end;
peercert(#proxy_socket{pp2_additional_info = AdditionalInfo}) ->
    proplists:get_value(pp2_ssl, AdditionalInfo, []).

%% @doc Peercert subject
-spec(peer_cert_subject(socket()) -> undefined | binary()).
peer_cert_subject(Sock) when is_port(Sock) ->
    undefined;
peer_cert_subject(#ssl_socket{ssl = SslSock}) ->
    case ssl:peercert(SslSock) of
        {ok, Cert} ->
            esockd_ssl:peer_cert_subject(Cert);
        _Error -> undefined
    end;
peer_cert_subject(Sock = #proxy_socket{}) ->
    %% Common Name? Haproxy PP2 will not pass subject.
    peer_cert_common_name(Sock).

%% @doc Peercert common name
-spec(peer_cert_common_name(socket()) -> undefined | binary()).
peer_cert_common_name(Sock) when is_port(Sock) ->
    undefined;
peer_cert_common_name(#ssl_socket{ssl = SslSock}) ->
    case ssl:peercert(SslSock) of
        {ok, Cert} ->
            esockd_ssl:peer_cert_common_name(Cert);
        _Error -> undefined
    end;
peer_cert_common_name(#proxy_socket{pp2_additional_info = AdditionalInfo}) ->
    proplists:get_value(pp2_ssl_cn,
                        proplists:get_value(pp2_ssl, AdditionalInfo, [])).

-spec(peersni(socket()) -> undefined | binary()).
peersni(Sock) when is_port(Sock) ->
    undefined;
peersni(#ssl_socket{ssl = SslSock}) ->
    case ssl:connection_information(SslSock, [sni_hostname]) of
        {ok, [{sni_hostname, SNI}]} when is_list(SNI) -> list_to_binary(SNI);
        %% If the client does not support SNI extensions,
        %% an empty list will be returned here
        %%
        %% see: https://github.com/erlang/otp/blob/e6ef9cb92499377c24d818376bfd60cbbcf68e60/lib/ssl/src/ssl.erl#L927-L935
        {ok, []} -> undefined;
        {error, _Reason} -> undefined
    end;
peersni(#proxy_socket{pp2_additional_info = AdditionalInfo}) ->
    proplists:get_value(pp2_authority, AdditionalInfo, undefined).

%% @doc Shutdown socket
-spec(shutdown(socket(), How) -> ok | {error, inet:posix()} when
    How :: read | write | read_write).
shutdown(Sock, How) when is_port(Sock) ->
    gen_tcp:shutdown(Sock, How);
shutdown(#ssl_socket{ssl = SslSock}, How) ->
    ssl:shutdown(SslSock, How);
shutdown(#proxy_socket{socket = Sock}, How) ->
    shutdown(Sock, How).

%% @doc TCP/DTLS socket -> #ssl_socket{}
-spec(ssl_upgrade_fun([ssl_option()]) -> esockd:sock_fun()).
ssl_upgrade_fun(SslOpts) ->
    {Timeout, SslOpts1} = take_handshake_timeout(SslOpts),
    {GCAfterHandshake, SslOpts2} = take_gc_after_handshake(SslOpts1),
    {fun ?MODULE:ssl_upgrade/3, [SslOpts2, #{timeout => Timeout,
                                             gc_after_handshake => GCAfterHandshake}]}.

ssl_upgrade(Sock, SslOpts1, #{timeout := Timeout,
                              gc_after_handshake := GCAfterHandshake}) ->
    try do_ssl_handshake(Sock, SslOpts1, Timeout) of
        {ok, NSock} ->
            GCAfterHandshake andalso gc(NSock),
            {ok, NSock};
        {error, Reason} when Reason =:= closed; Reason =:= timeout ->
            {error, Reason};
        {error, Reason} ->
            {error, {ssl_error, Reason}}
    catch
        _E:Reason ->
            {error, {ssl_failure, Reason}}
    end.

%% @private
do_ssl_handshake(Sock, SslOpts, Timeout) when is_port(Sock) ->
    case ssl:handshake(Sock, SslOpts, Timeout) of
        {ok, SslSock} ->
            {ok, #ssl_socket{tcp = Sock, ssl = SslSock}};
        {ok, SslSock, _Ext} -> %% OTP 21.0
            {ok, #ssl_socket{tcp = Sock, ssl = SslSock}};
        {error, _} = E -> E
    end;
%% dtls only
do_ssl_handshake(Sock, _SslOpts, Timeout) when is_record(Sock, sslsocket) ->
    case ssl:handshake(Sock, Timeout) of
        {ok, SslSock} ->
            {ok, #ssl_socket{tcp = undefined, ssl = SslSock}};
        {ok, SslSock, _Ext} -> %% OTP 21.0
            {ok, #ssl_socket{tcp = undefined, ssl = SslSock}};
        {error, _} = E -> E
    end;
do_ssl_handshake(ProxySock = #proxy_socket{socket = Sock}, SslOpts, Timeout) when is_port(Sock) ->
    case ssl:handshake(Sock, SslOpts, Timeout) of
        {ok, SslSock} ->
            {ok, ProxySock#proxy_socket{socket = #ssl_socket{tcp = Sock, ssl = SslSock}}};
        {ok, SslSock, _Ext} -> %% OTP 21.0
            {ok, ProxySock#proxy_socket{socket = #ssl_socket{tcp = Sock, ssl = SslSock}}};
        {error, _} = E -> E
    end.

take_handshake_timeout(SslOpts) ->
    case lists:keytake(handshake_timeout, 1, SslOpts) of
        {value, {handshake_timeout, Timeout}, SslOpts1} ->
            {Timeout, SslOpts1};
        false ->
            {?SSL_HANDSHAKE_TIMEOUT, SslOpts}
    end.

take_gc_after_handshake(SslOpts) ->
    {proplists:get_bool(gc_after_handshake, SslOpts),
     proplists:delete(gc_after_handshake, SslOpts)}.

%% @doc TCP | SSL -> ProxySocket
proxy_upgrade_fun(Opts) ->
    Timeout = proxy_protocol_timeout(Opts),
    {fun ?MODULE:proxy_upgrade/2, [Timeout]}.

proxy_upgrade(Sock, Timeout) ->
    case esockd_proxy_protocol:recv(?MODULE, Sock, Timeout) of
        {ok, ProxySock} -> {ok, ProxySock};
        {error, Reason} -> {error, Reason}
    end.

proxy_protocol_timeout(Opts) ->
    proplists:get_value(proxy_protocol_timeout, Opts, ?PROXY_RECV_TIMEOUT).

-spec(ensure_ok_or_exit(atom(), list(term())) -> term()).
ensure_ok_or_exit(Fun, Args = [Sock|_]) when is_atom(Fun), is_list(Args) ->
    case erlang:apply(?MODULE, Fun, Args) of
        {error, Reason} when Reason =:= enotconn; Reason =:= closed ->
            fast_close(Sock),
            exit(normal);
        {error, Reason} ->
            fast_close(Sock),
            exit({shutdown, Reason});
         Result -> Result
    end.

gc(Sock) when is_port(Sock) ->
    case erlang:port_info(Sock, connected) of
        {connected, Pid} ->
            erlang:garbage_collect(Pid),
            ok;
        undefined ->
            ok
    end;
%% Defined in ssl/src/ssl_api.hrl:
%% -record(sslsocket, {fd = nil, pid = nil}).
gc(#ssl_socket{ssl = {sslsocket, _, Pid}}) when is_pid(Pid) ->
    erlang:garbage_collect(Pid),
    ok;
%% In OTP 24+, the last element is a list of PIDs The first is spawned
%% by `ssl_gen_statem:init/1', the second by `tls_sender:init/1`.
gc(#ssl_socket{ssl = {sslsocket, _, Pids}}) when is_list(Pids) ->
    lists:foreach(
      fun(Pid) when is_pid(Pid) ->
              erlang:garbage_collect(Pid);
         (_) ->
              ok
      end,
      Pids);
gc(#proxy_socket{socket = Sock}) ->
    gc(Sock);
gc(_Sock) -> ok.

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_transport).
18
19	-include("esockd.hrl").
20	-include_lib("ssl/src/ssl_api.hrl").
21
22	-export([type/1, is_ssl/1]).
23	-export([listen/2]).
24	-export([ready/3, wait/1]).
25	-export([send/2, async_send/2, async_send/3, recv/2, recv/3, async_recv/2, async_recv/3]).
26	-export([controlling_process/2]).
27	-export([close/1, fast_close/1]).
28	-export([getopts/2, setopts/2, getstat/2]).
29	-export([sockname/1, peername/1, shutdown/2]).
30	-export([peercert/1, peer_cert_subject/1, peer_cert_common_name/1, peersni/1]).
31	-export([ensure_ok_or_exit/2]).
32	-export([gc/1]).
33
34	%% Internal callbacks
35	-export([ ssl_upgrade_fun/1
36	, ssl_upgrade/3
37	, proxy_upgrade_fun/1
38	, proxy_upgrade/2
39	]).
40
41	-export_type([socket/0]).
42
43	-type(ssl_socket() :: #ssl_socket{}).
44	-type(proxy_socket() :: #proxy_socket{}).
45	-type(socket() :: inet:socket() \| ssl_socket() \| proxy_socket() \| #sslsocket{}).
46
47	-spec(type(socket()) -> tcp \| ssl \| proxy).
48	type(Sock) when is_port(Sock) ->
49	tcp;	1✔
50	type(#ssl_socket{}) ->
51	ssl;	1✔
52	type(#proxy_socket{}) ->
53	proxy.	1✔
54
55	-spec(is_ssl(socket()) -> boolean()).
56	is_ssl(Sock) when is_port(Sock) ->
57	false;	2✔
58	is_ssl(#ssl_socket{}) ->
59	true;	2✔
60	is_ssl(#proxy_socket{socket = Sock}) ->
61	is_ssl(Sock).	2✔
62
63	-spec(ready(pid(), socket(), [esockd:sock_fun()]) -> any()).
64	ready(Pid, Sock, UpgradeFuns) ->
65	Pid ! {sock_ready, Sock, UpgradeFuns}.	44✔
66
67	-spec(wait(socket()) -> {ok, socket()} \| {error, term()}).
68	wait(Sock) ->
69	receive	44✔
70	{sock_ready, Sock, UpgradeFuns} ->
71	upgrade(Sock, UpgradeFuns)	44✔
72	end.
73
74	-spec(upgrade(socket(), [esockd:sock_fun()]) -> {ok, socket()} \| {error, term()}).
75	upgrade(Sock, []) ->
76	{ok, Sock};	40✔
77	upgrade(Sock, [{Fun, Args}\|More]) ->
78	case apply(Fun, [Sock\|Args]) of	25✔
79	{ok, NewSock} -> upgrade(NewSock, More);	21✔
80	Error -> fast_close(Sock), Error	4✔
81	end.
82
83	-spec(listen(inet:port_number(), [gen_tcp:listen_option()])
84	-> {ok, inet:socket()} \| {error, system_limit \| inet:posix()}).
85	listen(Port, Opts) ->
86	gen_tcp:listen(Port, Opts).	57✔
87
88	-spec(controlling_process(socket(), pid()) -> ok \| {error, Reason} when
89	Reason :: closed \| not_owner \| badarg \| inet:posix()).
90	controlling_process(Sock, NewOwner) when is_port(Sock) ->
91	gen_tcp:controlling_process(Sock, NewOwner);	38✔
92	controlling_process(#ssl_socket{ssl = SslSock}, NewOwner) ->
93	ssl:controlling_process(SslSock, NewOwner);	×
94	controlling_process(#proxy_socket{socket = Sock}, NewOwner) ->
95	controlling_process(Sock, NewOwner);	×
96	%% Before upgrading, the DTLS socket is #sslsocket{}
97	%% type, instead of a port().
98	%%
99	%% See: ssl:transport_accept/1
100	controlling_process(SslSock = #sslsocket{}, NewOwner) ->
101	ssl:controlling_process(SslSock, NewOwner).	6✔
102
103	-spec(close(socket()) -> ok \| {error, term()}).
104	close(Sock) when is_port(Sock) ->
105	gen_tcp:close(Sock);	19✔
106	close(#ssl_socket{ssl = SslSock}) ->
107	ssl:close(SslSock);	3✔
108	close(#proxy_socket{socket = Sock}) ->
109	close(Sock).	1✔
110
111	-spec(fast_close(socket()) -> ok).
112	fast_close(Sock) when is_port(Sock) ->
113	catch port_close(Sock), ok;	50✔
114	fast_close(#ssl_socket{tcp = Sock, ssl = SslSock}) ->
115	_ = fast_close_sslsock(SslSock),	×
116	catch port_close(Sock), ok;	×
117	fast_close(SslSock = #sslsocket{}) ->
118	fast_close_sslsock(SslSock);	10✔
119	fast_close(#proxy_socket{socket = Sock}) ->
120	fast_close(Sock).	×
121
122	%% @private
123	fast_close_sslsock(SslSock) ->
124	{Pid, MRef} = spawn_monitor(fun() -> ssl:close(SslSock) end),	10✔
125	TRef = erlang:send_after(?SSL_CLOSE_TIMEOUT, self(), {Pid, ssl_close_timeout}),	10✔
126	_ = receive	10✔
127	{Pid, ssl_close_timeout} ->
128	erlang:demonitor(MRef, [flush]),	×
129	exit(Pid, kill);	×
130	{'DOWN', MRef, process, Pid, _Reason} ->
131	erlang:cancel_timer(TRef)	10✔
132	end,
133	ok.	10✔
134
135	-spec(send(socket(), iodata()) -> ok \| {error, Reason} when
136	Reason :: closed \| timeout \| inet:posix()).
137	send(Sock, Data) when is_port(Sock) ->
138	gen_tcp:send(Sock, Data);	17✔
139	send(#ssl_socket{ssl = SslSock}, Data) ->
140	ssl:send(SslSock, Data);	6✔
141	send(#proxy_socket{socket = Sock}, Data) ->
142	send(Sock, Data).	8✔
143
144	%% @doc Port command to write data.
145	-spec(async_send(socket(), iodata()) -> ok \| {error, Reason} when
146	Reason :: close \| timeout \| inet:posix()).
147	async_send(Sock, Data) ->
148	async_send(Sock, Data, []).	6✔
149
150	-spec(async_send(socket(), iodata(), [force\|nosuspend]) -> ok \| {error, Reason} when
151	Reason :: close \| timeout \| inet:posix()).
152	async_send(Sock, Data, Options) when is_port(Sock) ->
153	do_port_command(Sock, Data, Options);	3✔
154	async_send(Sock = #ssl_socket{ssl = SslSock}, Data, _Options) ->
155	case ssl:send(SslSock, Data) of	3✔
156	ok ->
157	%% TODO: get rid of this message send
158	self() ! {inet_reply, Sock, ok},	3✔
159	ok;	3✔
160	Error ->
161	Error	×
162	end;
163	async_send(#proxy_socket{socket = Sock}, Data, _Options) ->
164	async_send(Sock, Data).	×
165
166	-if(?OTP_RELEASE >= 26).
167	do_port_command(Sock, Data, _Options) ->
168	case gen_tcp:send(Sock, Data) of
169	ok ->
170	%% TODO: get rid of this message send
171	self() ! {inet_reply, Sock, ok},
172	ok;
173	{error, Reason} ->
174	{error, Reason}
175	end.
176	-else.
177	do_port_command(Sock, Data, Options) ->
178	try erlang:port_command(Sock, Data, Options) of	3✔
179	true ->
180	ok;	3✔
181	false -> %% nosuspend option and port busy
UNCOV 182	{error, busy}	×
183	catch
184	error:_Error ->
UNCOV 185	{error, einval}	×
186	end.
187	-endif.
188
189	-spec(recv(socket(), non_neg_integer())
190	-> {ok, iodata()} \| {error, closed \| inet:posix()}).
191	recv(Sock, Length) when is_port(Sock) ->
192	gen_tcp:recv(Sock, Length);	46✔
193	recv(#ssl_socket{ssl = SslSock}, Length) ->
194	ssl:recv(SslSock, Length);	9✔
195	recv(#proxy_socket{socket = Sock}, Length) ->
196	recv(Sock, Length).	16✔
197
198	-spec(recv(socket(), non_neg_integer(), timeout())
199	-> {ok, iodata()} \| {error, closed \| inet:posix()}).
200	recv(Sock, Length, Timeout) when is_port(Sock) ->
201	gen_tcp:recv(Sock, Length, Timeout);	13✔
202	recv(#ssl_socket{ssl = SslSock}, Length, Timeout) ->
203	ssl:recv(SslSock, Length, Timeout);	×
204	recv(#proxy_socket{socket = Sock}, Length, Timeout) ->
205	recv(Sock, Length, Timeout).	×
206
207	%% @doc Async receive data.
208	-spec(async_recv(socket(), non_neg_integer()) -> {ok, reference()}).
209	async_recv(Sock, Length) ->
210	async_recv(Sock, Length, infinity).	3✔
211
212	-spec(async_recv(socket(), non_neg_integer(), timeout()) -> {ok, reference()}).
213	async_recv(Sock = #ssl_socket{ssl = SslSock}, Length, Timeout) ->
214	Self = self(),	×
215	Ref = make_ref(),	×
216	spawn(fun() ->	×
217	Self ! {inet_async, Sock, Ref, ssl:recv(SslSock, Length, Timeout)}	×
218	end),
219	{ok, Ref};	×
220	async_recv(Sock, Length, infinity) when is_port(Sock) ->
221	prim_inet:async_recv(Sock, Length, -1);	3✔
222	async_recv(Sock, Length, Timeout) when is_port(Sock) ->
223	prim_inet:async_recv(Sock, Length, Timeout);	×
224	async_recv(#proxy_socket{socket = Sock}, Length, Timeout) ->
225	async_recv(Sock, Length, Timeout).	×
226
227	%% @doc Get socket options.
228	-spec(getopts(socket(), [inet:socket_getopt()])
229	-> {ok, [inet:socket_setopt()]} \| {error, inet:posix()}).
230	getopts(Sock, OptionNames) when is_port(Sock) ->
231	inet:getopts(Sock, OptionNames);	14✔
232	getopts(#ssl_socket{ssl = SslSock}, OptionNames) ->
233	ssl:getopts(SslSock, OptionNames);	×
234	getopts(SslSock = #sslsocket{}, OptionNames) ->
235	ssl:getopts(SslSock, OptionNames);	×
236	getopts(#proxy_socket{socket = Sock}, OptionNames) ->
237	getopts(Sock, OptionNames).	×
238
239	%% @doc Set socket options
240	-spec(setopts(socket(), [inet:socket_setopt()]) -> ok \| {error, inet:posix()}).
241	setopts(Sock, Opts) when is_port(Sock) ->
242	inet:setopts(Sock, Opts);	30✔
243	setopts(#ssl_socket{ssl = SslSock}, Opts) ->
244	ssl:setopts(SslSock, Opts);	×
245	setopts(SslSock = #sslsocket{}, Opts) ->
246	ssl:setopts(SslSock, Opts);	×
247	setopts(#proxy_socket{socket = Socket}, Opts) ->
248	setopts(Socket, Opts).	×
249
250	%% @doc Get socket stats
251	-spec(getstat(socket(), [inet:stat_option()])
252	-> {ok, [{inet:stat_option(), integer()}]} \| {error, inet:posix()}).
253	getstat(Sock, Stats) when is_port(Sock) ->
254	inet:getstat(Sock, Stats);	1✔
255	getstat(#ssl_socket{tcp = Sock}, Stats) when is_port(Sock) ->
256	inet:getstat(Sock, Stats);	×
257	getstat(#ssl_socket{ssl = SslSock}, Stats) ->
258	ssl_getstat(SslSock, Stats);	2✔
259	getstat(SslSock = #sslsocket{}, Stats) ->
260	ssl_getstat(SslSock, Stats);	×
261	getstat(#proxy_socket{socket = Sock}, Stats) ->
262	getstat(Sock, Stats).	×
263
264	%% Below OTP 23, the ssl:getstats/2 will get a funcation_clause
265	%% exception for dtls socket
266	-if(?OTP_RELEASE >= 23).
267	ssl_getstat(SslSock, Stats) ->
268	ssl:getstat(SslSock, Stats).	2✔
269	-else.
270	ssl_getstat(#sslsocket{fd = {gen_udp, {_,{{_sockhost,_sockport}, SockPort}},_}}, Stats) ->
271	inet:getstat(SockPort, Stats);
272	ssl_getstat(SslSock, Stats) ->
273	ssl:getstat(SslSock, Stats).
274	-endif.
275
276	%% @doc Sockname
277	-spec(sockname(socket()) -> {ok, {inet:ip_address(), inet:port_number()}} \|
278	{error, inet:posix()}).
279	sockname(Sock) when is_port(Sock) ->
280	inet:sockname(Sock);	3✔
281	sockname(#ssl_socket{ssl = SslSock}) ->
282	ssl:sockname(SslSock);	×
283	sockname(SslSock = #sslsocket{}) ->
284	ssl:sockname(SslSock);	×
285	sockname(#proxy_socket{dst_addr = DstAddr, dst_port = DstPort}) ->
286	{ok, {DstAddr, DstPort}}.	×
287
288	%% @doc Peername
289	-spec(peername(socket()) -> {ok, {inet:ip_address(), inet:port_number()}} \|
290	{error, inet:posix()}).
291	peername(Sock) when is_port(Sock) ->
292	inet:peername(Sock);	39✔
293	peername(#ssl_socket{ssl = SslSock}) ->
294	ssl:peername(SslSock);	2✔
295	peername(#proxy_socket{src_addr = SrcAddr, src_port = SrcPort}) ->
296	{ok, {SrcAddr, SrcPort}};	×
297	%% Before upgrading, the DTLS socket is #sslsocket{}
298	%% type, instead of a port().
299	%%
300	%% See: ssl:transport_accept/1
301	peername(SslSock = #sslsocket{}) ->
302	ssl:peername(SslSock).	6✔
303
304	%% @doc Socket peercert
305	-spec(peercert(socket()) -> nossl \| binary() \| list(pp2_additional_ssl_field()) \|
306	{error, term()}).
307	peercert(Sock) when is_port(Sock) ->
308	nossl;	1✔
309	peercert(#ssl_socket{ssl = SslSock}) ->
310	case ssl:peercert(SslSock) of	×
311	{ok, Cert} -> Cert;	×
312	%% One-way SSL
313	{error, no_peercert} ->
314	undefined;	×
315	Error -> Error	×
316	end;
317	peercert(SslSock = #sslsocket{}) ->
318	case ssl:peercert(SslSock) of	×
319	{ok, Cert} -> Cert;	×
320	{error, no_peercert} -> undefined;	×
321	Error -> Error	×
322	end;
323	peercert(#proxy_socket{pp2_additional_info = AdditionalInfo}) ->
324	proplists:get_value(pp2_ssl, AdditionalInfo, []).	×
325
326	%% @doc Peercert subject
327	-spec(peer_cert_subject(socket()) -> undefined \| binary()).
328	peer_cert_subject(Sock) when is_port(Sock) ->
329	undefined;	1✔
330	peer_cert_subject(#ssl_socket{ssl = SslSock}) ->
331	case ssl:peercert(SslSock) of	×
332	{ok, Cert} ->
333	esockd_ssl:peer_cert_subject(Cert);	×
334	_Error -> undefined	×
335	end;
336	peer_cert_subject(Sock = #proxy_socket{}) ->
337	%% Common Name? Haproxy PP2 will not pass subject.
338	peer_cert_common_name(Sock).	×
339
340	%% @doc Peercert common name
341	-spec(peer_cert_common_name(socket()) -> undefined \| binary()).
342	peer_cert_common_name(Sock) when is_port(Sock) ->
343	undefined;	1✔
344	peer_cert_common_name(#ssl_socket{ssl = SslSock}) ->
345	case ssl:peercert(SslSock) of	×
346	{ok, Cert} ->
347	esockd_ssl:peer_cert_common_name(Cert);	×
348	_Error -> undefined	×
349	end;
350	peer_cert_common_name(#proxy_socket{pp2_additional_info = AdditionalInfo}) ->
351	proplists:get_value(pp2_ssl_cn,	×
352	proplists:get_value(pp2_ssl, AdditionalInfo, [])).
353
354	-spec(peersni(socket()) -> undefined \| binary()).
355	peersni(Sock) when is_port(Sock) ->
356	undefined;	1✔
357	peersni(#ssl_socket{ssl = SslSock}) ->
358	case ssl:connection_information(SslSock, [sni_hostname]) of	2✔
359	{ok, [{sni_hostname, SNI}]} when is_list(SNI) -> list_to_binary(SNI);	1✔
360	%% If the client does not support SNI extensions,
361	%% an empty list will be returned here
362	%%
363	%% see: https://github.com/erlang/otp/blob/e6ef9cb92499377c24d818376bfd60cbbcf68e60/lib/ssl/src/ssl.erl#L927-L935
364	{ok, []} -> undefined;	1✔
365	{error, _Reason} -> undefined	×
366	end;
367	peersni(#proxy_socket{pp2_additional_info = AdditionalInfo}) ->
368	proplists:get_value(pp2_authority, AdditionalInfo, undefined).	1✔
369
370	%% @doc Shutdown socket
371	-spec(shutdown(socket(), How) -> ok \| {error, inet:posix()} when
372	How :: read \| write \| read_write).
373	shutdown(Sock, How) when is_port(Sock) ->
374	gen_tcp:shutdown(Sock, How);	1✔
375	shutdown(#ssl_socket{ssl = SslSock}, How) ->
376	ssl:shutdown(SslSock, How);	×
377	shutdown(#proxy_socket{socket = Sock}, How) ->
378	shutdown(Sock, How).	×
379
380	%% @doc TCP/DTLS socket -> #ssl_socket{}
381	-spec(ssl_upgrade_fun([ssl_option()]) -> esockd:sock_fun()).
382	ssl_upgrade_fun(SslOpts) ->
383	{Timeout, SslOpts1} = take_handshake_timeout(SslOpts),	11✔
384	{GCAfterHandshake, SslOpts2} = take_gc_after_handshake(SslOpts1),	11✔
385	{fun ?MODULE:ssl_upgrade/3, [SslOpts2, #{timeout => Timeout,	11✔
386	gc_after_handshake => GCAfterHandshake}]}.
387
388	ssl_upgrade(Sock, SslOpts1, #{timeout := Timeout,
389	gc_after_handshake := GCAfterHandshake}) ->
390	try do_ssl_handshake(Sock, SslOpts1, Timeout) of	11✔
391	{ok, NSock} ->
392	GCAfterHandshake andalso gc(NSock),	11✔
393	{ok, NSock};	11✔
394	{error, Reason} when Reason =:= closed; Reason =:= timeout ->
395	{error, Reason};	×
396	{error, Reason} ->
397	{error, {ssl_error, Reason}}	×
398	catch
399	_E:Reason ->
400	{error, {ssl_failure, Reason}}	×
401	end.
402
403	%% @private
404	do_ssl_handshake(Sock, SslOpts, Timeout) when is_port(Sock) ->
405	case ssl:handshake(Sock, SslOpts, Timeout) of	5✔
406	{ok, SslSock} ->
407	{ok, #ssl_socket{tcp = Sock, ssl = SslSock}};	5✔
408	{ok, SslSock, _Ext} -> %% OTP 21.0
409	{ok, #ssl_socket{tcp = Sock, ssl = SslSock}};	×
410	{error, _} = E -> E	×
411	end;
412	%% dtls only
413	do_ssl_handshake(Sock, _SslOpts, Timeout) when is_record(Sock, sslsocket) ->
414	case ssl:handshake(Sock, Timeout) of	6✔
415	{ok, SslSock} ->
416	{ok, #ssl_socket{tcp = undefined, ssl = SslSock}};	6✔
417	{ok, SslSock, _Ext} -> %% OTP 21.0
418	{ok, #ssl_socket{tcp = undefined, ssl = SslSock}};	×
419	{error, _} = E -> E	×
420	end;
421	do_ssl_handshake(ProxySock = #proxy_socket{socket = Sock}, SslOpts, Timeout) when is_port(Sock) ->
422	case ssl:handshake(Sock, SslOpts, Timeout) of	×
423	{ok, SslSock} ->
424	{ok, ProxySock#proxy_socket{socket = #ssl_socket{tcp = Sock, ssl = SslSock}}};	×
425	{ok, SslSock, _Ext} -> %% OTP 21.0
426	{ok, ProxySock#proxy_socket{socket = #ssl_socket{tcp = Sock, ssl = SslSock}}};	×
427	{error, _} = E -> E	×
428	end.
429
430	take_handshake_timeout(SslOpts) ->
431	case lists:keytake(handshake_timeout, 1, SslOpts) of	11✔
432	{value, {handshake_timeout, Timeout}, SslOpts1} ->
433	{Timeout, SslOpts1};	2✔
434	false ->
435	{?SSL_HANDSHAKE_TIMEOUT, SslOpts}	9✔
436	end.
437
438	take_gc_after_handshake(SslOpts) ->
439	{proplists:get_bool(gc_after_handshake, SslOpts),	11✔
440	proplists:delete(gc_after_handshake, SslOpts)}.
441
442	%% @doc TCP \| SSL -> ProxySocket
443	proxy_upgrade_fun(Opts) ->
444	Timeout = proxy_protocol_timeout(Opts),	14✔
445	{fun ?MODULE:proxy_upgrade/2, [Timeout]}.	14✔
446
447	proxy_upgrade(Sock, Timeout) ->
448	case esockd_proxy_protocol:recv(?MODULE, Sock, Timeout) of	13✔
449	{ok, ProxySock} -> {ok, ProxySock};	9✔
450	{error, Reason} -> {error, Reason}	4✔
451	end.
452
453	proxy_protocol_timeout(Opts) ->
454	proplists:get_value(proxy_protocol_timeout, Opts, ?PROXY_RECV_TIMEOUT).	14✔
455
456	-spec(ensure_ok_or_exit(atom(), list(term())) -> term()).
457	ensure_ok_or_exit(Fun, Args = [Sock\|_]) when is_atom(Fun), is_list(Args) ->
458	case erlang:apply(?MODULE, Fun, Args) of	2✔
459	{error, Reason} when Reason =:= enotconn; Reason =:= closed ->
460	fast_close(Sock),	×
461	exit(normal);	×
462	{error, Reason} ->
463	fast_close(Sock),	×
464	exit({shutdown, Reason});	×
465	Result -> Result	2✔
466	end.
467
468	gc(Sock) when is_port(Sock) ->
469	case erlang:port_info(Sock, connected) of	1✔
470	{connected, Pid} ->
471	erlang:garbage_collect(Pid),	1✔
472	ok;	1✔
473	undefined ->
474	ok	×
475	end;
476	%% Defined in ssl/src/ssl_api.hrl:
477	%% -record(sslsocket, {fd = nil, pid = nil}).
478	gc(#ssl_socket{ssl = {sslsocket, _, Pid}}) when is_pid(Pid) ->
479	erlang:garbage_collect(Pid),	×
480	ok;	×
481	%% In OTP 24+, the last element is a list of PIDs The first is spawned
482	%% by `ssl_gen_statem:init/1', the second by `tls_sender:init/1`.
483	gc(#ssl_socket{ssl = {sslsocket, _, Pids}}) when is_list(Pids) ->
484	lists:foreach(	4✔
485	fun(Pid) when is_pid(Pid) ->
486	erlang:garbage_collect(Pid);	8✔
487	(_) ->
488	ok	×
489	end,
490	Pids);
491	gc(#proxy_socket{socket = Sock}) ->
492	gc(Sock);	×
493	gc(_Sock) -> ok.	×

emqx / esockd / 349

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