emqx / esockd / 372

Committed 12 Mar 2024 08:41AM UTC coverage: 67.341% (+0.2%) from 67.117%

Build # 372

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #186 from emqx/fast_fail_on_invalid_ssl_opts

fix: fast fail on invalid ssl options

Run Details

2 of 2 new or added lines in 1 file covered. (100.0%)

1 existing line in 1 file now uncovered.

699 of 1038 relevant lines covered (67.34%)

160.22 hits per line

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

50.31

/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) ->
    try erlang:port_command(Sock, Data, Options) of
        true -> ok;
        false -> %% nosuspend option and port busy
            {error, busy}
    catch
        error:_Error ->
            {error, einval}
    end;
async_send(Sock = #ssl_socket{ssl = SslSock}, Data, _Options) ->
    case ssl:send(SslSock, Data) of
        ok -> self() ! {inet_reply, Sock, ok}, ok;
        Error -> Error
    end;
async_send(#proxy_socket{socket = Sock}, Data, _Options) ->
    async_send(Sock, Data).

-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: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}]}.

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

emqx / esockd / 372

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