<commit_before>%%% @doc Connects to the server and performs all the necessary handshakes.
%%%
%%% Special kind of command - it's exclusive: no other commands can run until
%%% this one finishes.
%%% It also uses some "private" epgsql_sock's APIs
%%%
-module(epgsql_cmd_connect).
-behaviour(epgsql_command).
-export([hide_password/1, opts_hide_password/1, open_socket/2]).
-export([init/1, execute/2, handle_message/4]).
-export_type([response/0, connect_error/0]).

-type response() :: connected
                  | {error, connect_error()}.
-type connect_error() ::
        invalid_authorization_specification
      | invalid_password
      | {unsupported_auth_method,
         kerberosV5 | crypt | scm | gss | sspi | {unknown, integer()} | {sasl, [binary()]}}
      | {sasl_server_final, any()}
      | epgsql:query_error().

-include("epgsql.hrl").
-include("protocol.hrl").

-type auth_fun() :: fun((init | binary(), _, _) ->
                                     {send, byte(), iodata(), any()}
                                   | ok
                                   | {error, any()}
                                   | unknown).

-record(connect,
        {opts :: map(),
         auth_fun :: auth_fun() | undefined,
         auth_state :: any() | undefined,
         auth_send :: {integer(), iodata()} | undefined,
         stage = connect :: connect | maybe_auth | auth | initialization}).

-define(SCRAM_AUTH_METHOD, <<"SCRAM-SHA-256">>).
-define(AUTH_OK, 0).
-define(AUTH_CLEARTEXT, 3).
-define(AUTH_MD5, 5).
-define(AUTH_SASL, 10).
-define(AUTH_SASL_CONTINUE, 11).
-define(AUTH_SASL_FINAL, 12).

init(#{host := _, username := _} = Opts) ->
    #connect{opts = Opts}.

execute(PgSock, #connect{opts = #{username := Username} = Opts, stage = connect} = State) ->
    SockOpts = [{active, false}, {packet, raw}, binary, {nodelay, true}, {keepalive, true}],
    PgSock1 = epgsql_sock:set_attr(connect_opts, Opts, PgSock),
    case open_socket(SockOpts, Opts) of
        {ok, Mode, Sock} ->
            PgSock2 = epgsql_sock:set_net_socket(Mode, Sock, PgSock1),
            Opts2 = ["user", 0, Username, 0],
            Opts3 = case maps:find(database, Opts) of
                        error -> Opts2;
                        {ok, Database}  -> [Opts2 | ["database", 0, Database, 0]]
                    end,
           {Opts4, PgSock3} =
               case Opts of
                   #{replication := Replication}  ->
                       {[Opts3 | ["replication", 0, Replication, 0]],
                        epgsql_sock:init_replication_state(PgSock2)};
                   _ -> {Opts3, PgSock2}
               end,
            Opts5 = case Opts of
                        #{application_name := ApplicationName}  ->
                            [Opts4 | ["application_name", 0, ApplicationName, 0]];
                        _ ->
                            Opts4
                    end,
           ok = epgsql_sock:send(PgSock2, [<<196608:?int32>>, Opts5, 0]),
           PgSock4 = case Opts of
                         #{async := Async} ->
                             epgsql_sock:set_attr(async, Async, PgSock3);
                         _ -> PgSock3
                     end,
           {ok, PgSock4, State#connect{stage = maybe_auth}};
        {error, Reason} = Error ->
            {stop, Reason, Error, PgSock}
    end;
execute(PgSock, #connect{stage = auth, auth_send = {PacketId, Data}} = St) ->
    ok = epgsql_sock:send(PgSock, PacketId, Data),
    {ok, PgSock, St#connect{auth_send = undefined}}.

-spec open_socket([{atom(), any()}], epgsql:connect_opts()) ->
    {ok , gen_tcp | ssl, port() | ssl:sslsocket()} | {error, any()}.
open_socket(SockOpts, #{host := Host} = ConnectOpts) ->
    Timeout = maps:get(timeout, ConnectOpts, 5000),
    Deadline = deadline(Timeout),
    Port = maps:get(port, ConnectOpts, 5432),
    case gen_tcp:connect(Host, Port, SockOpts, Timeout) of
       {ok, Sock} ->
           client_handshake(Sock, ConnectOpts, Deadline);
       {error, _Reason} = Error ->
           Error
    end.

client_handshake(Sock, ConnectOpts, Deadline) ->
    %% Increase the buffer size.  Following the recommendation in the inet man page:
    %%
    %%    It is recommended to have val(buffer) >=
    %%    max(val(sndbuf),val(recbuf)).

    {ok, [{recbuf, RecBufSize}, {sndbuf, SndBufSize}]} =
        inet:getopts(Sock, [recbuf, sndbuf]),
    inet:setopts(Sock, [{buffer, max(RecBufSize, SndBufSize)}]),
    maybe_ssl(Sock, maps:get(ssl, ConnectOpts, false), ConnectOpts, Deadline).

maybe_ssl(Sock, false, _ConnectOpts, _Deadline) ->
    {ok, gen_tcp, Sock};
maybe_ssl(Sock, Flag, ConnectOpts, Deadline) ->
    ok = gen_tcp:send(Sock, <<8:?int32, 80877103:?int32>>),
    Timeout0 = timeout(Deadline),
    case gen_tcp:recv(Sock, 1, Timeout0) of
        {ok, <<$S>>}  ->
            SslOpts = maps:get(ssl_opts, ConnectOpts, []),
            Timeout = timeout(Deadline),
            case ssl:connect(Sock, SslOpts, Timeout) of
                {ok, Sock2} ->
                    {ok, ssl, Sock2};
                {error, Reason} ->
                    Err = {ssl_negotiation_failed, Reason},
                    {error, Err}
            end;
        {ok, <<$N>>} ->
            case Flag of
                true ->
                   {ok, gen_tcp, Sock};
                required ->
                    {error, ssl_not_available}
            end;
        {error, Reason} ->
            {error, Reason}
    end.

%% @doc Replace `password' in Opts map with obfuscated one
opts_hide_password(#{password := Password} = Opts) ->
    HiddenPassword = hide_password(Password),
    Opts#{password => HiddenPassword};
opts_hide_password(Opts) -> Opts.


%% @doc this function wraps plaintext password to a lambda function, so, if
%% epgsql_sock process crashes when executing `connect' command, password will<commit_msg>Filter sensitive info from connection_opts prior to setting it into the epgsql_sock state<commit_after>%%% @doc Connects to the server and performs all the necessary handshakes.
%%%
%%% Special kind of command - it's exclusive: no other commands can run until
%%% this one finishes.
%%% It also uses some "private" epgsql_sock's APIs
%%%
-module(epgsql_cmd_connect).
-behaviour(epgsql_command).
-export([hide_password/1, opts_hide_password/1, open_socket/2]).
-export([init/1, execute/2, handle_message/4]).
-export_type([response/0, connect_error/0]).

-type response() :: connected
                  | {error, connect_error()}.
-type connect_error() ::
        invalid_authorization_specification
      | invalid_password
      | {unsupported_auth_method,
         kerberosV5 | crypt | scm | gss | sspi | {unknown, integer()} | {sasl, [binary()]}}
      | {sasl_server_final, any()}
      | epgsql:query_error().

-include("epgsql.hrl").
-include("protocol.hrl").

-type auth_fun() :: fun((init | binary(), _, _) ->
                                     {send, byte(), iodata(), any()}
                                   | ok
                                   | {error, any()}
                                   | unknown).

-record(connect,
        {opts :: map(),
         auth_fun :: auth_fun() | undefined,
         auth_state :: any() | undefined,
         auth_send :: {integer(), iodata()} | undefined,
         stage = connect :: connect | maybe_auth | auth | initialization}).

-define(SCRAM_AUTH_METHOD, <<"SCRAM-SHA-256">>).
-define(AUTH_OK, 0).
-define(AUTH_CLEARTEXT, 3).
-define(AUTH_MD5, 5).
-define(AUTH_SASL, 10).
-define(AUTH_SASL_CONTINUE, 11).
-define(AUTH_SASL_FINAL, 12).

init(#{host := _, username := _} = Opts) ->
    #connect{opts = Opts}.

execute(PgSock, #connect{opts = #{username := Username} = Opts, stage = connect} = State) ->
    SockOpts = [{active, false}, {packet, raw}, binary, {nodelay, true}, {keepalive, true}],
    FilteredOpts = filter_sensitive_info(Opts),
    PgSock1 = epgsql_sock:set_attr(connect_opts, FilteredOpts, PgSock),
    case open_socket(SockOpts, Opts) of
        {ok, Mode, Sock} ->
            PgSock2 = epgsql_sock:set_net_socket(Mode, Sock, PgSock1),
            Opts2 = ["user", 0, Username, 0],
            Opts3 = case maps:find(database, Opts) of
                        error -> Opts2;
                        {ok, Database}  -> [Opts2 | ["database", 0, Database, 0]]
                    end,
           {Opts4, PgSock3} =
               case Opts of
                   #{replication := Replication}  ->
                       {[Opts3 | ["replication", 0, Replication, 0]],
                        epgsql_sock:init_replication_state(PgSock2)};
                   _ -> {Opts3, PgSock2}
               end,
            Opts5 = case Opts of
                        #{application_name := ApplicationName}  ->
                            [Opts4 | ["application_name", 0, ApplicationName, 0]];
                        _ ->
                            Opts4
                    end,
           ok = epgsql_sock:send(PgSock2, [<<196608:?int32>>, Opts5, 0]),
           PgSock4 = case Opts of
                         #{async := Async} ->
                             epgsql_sock:set_attr(async, Async, PgSock3);
                         _ -> PgSock3
                     end,
           {ok, PgSock4, State#connect{stage = maybe_auth}};
        {error, Reason} = Error ->
            {stop, Reason, Error, PgSock}
    end;
execute(PgSock, #connect{stage = auth, auth_send = {PacketId, Data}} = St) ->
    ok = epgsql_sock:send(PgSock, PacketId, Data),
    {ok, PgSock, St#connect{auth_send = undefined}}.

-spec open_socket([{atom(), any()}], epgsql:connect_opts()) ->
    {ok , gen_tcp | ssl, port() | ssl:sslsocket()} | {error, any()}.
open_socket(SockOpts, #{host := Host} = ConnectOpts) ->
    Timeout = maps:get(timeout, ConnectOpts, 5000),
    Deadline = deadline(Timeout),
    Port = maps:get(port, ConnectOpts, 5432),
    case gen_tcp:connect(Host, Port, SockOpts, Timeout) of
       {ok, Sock} ->
           client_handshake(Sock, ConnectOpts, Deadline);
       {error, _Reason} = Error ->
           Error
    end.

client_handshake(Sock, ConnectOpts, Deadline) ->
    %% Increase the buffer size.  Following the recommendation in the inet man page:
    %%
    %%    It is recommended to have val(buffer) >=
    %%    max(val(sndbuf),val(recbuf)).

    {ok, [{recbuf, RecBufSize}, {sndbuf, SndBufSize}]} =
        inet:getopts(Sock, [recbuf, sndbuf]),
    inet:setopts(Sock, [{buffer, max(RecBufSize, SndBufSize)}]),
    maybe_ssl(Sock, maps:get(ssl, ConnectOpts, false), ConnectOpts, Deadline).

maybe_ssl(Sock, false, _ConnectOpts, _Deadline) ->
    {ok, gen_tcp, Sock};
maybe_ssl(Sock, Flag, ConnectOpts, Deadline) ->
    ok = gen_tcp:send(Sock, <<8:?int32, 80877103:?int32>>),
    Timeout0 = timeout(Deadline),
    case gen_tcp:recv(Sock, 1, Timeout0) of
        {ok, <<$S>>}  ->
            SslOpts = maps:get(ssl_opts, ConnectOpts, []),
            Timeout = timeout(Deadline),
            case ssl:connect(Sock, SslOpts, Timeout) of
                {ok, Sock2} ->
                    {ok, ssl, Sock2};
                {error, Reason} ->
                    Err = {ssl_negotiation_failed, Reason},
                    {error, Err}
            end;
        {ok, <<$N>>} ->
            case Flag of
                true ->
                   {ok, gen_tcp, Sock};
                required ->
                    {error, ssl_not_available}
            end;
        {error, Reason} ->
            {error, Reason}
    end.

%% @doc Replace `password' in Opts map with obfuscated one
opts_hide_password(#{password := Password} = Opts) ->
    HiddenPassword = hide_password(Password),
    Opts#{password => <PASSWORD>};
opts_hide_password(Opts) -> Opts.

%% @doc password and username are sensitive data that should not be stored in a
%% permanent state that might crash during code upgrade
filter_sensitive_info(Opts0) ->
  maps:without([password, username], Opts0).

%% @doc this function wraps plaintext password to a lambda function, so, if
%% epgsql_sock process crashes when executing `connect' command, password will
{'fraction_non_alphanumeric': 0.14060293995515322, 'fraction_numerical': 0.014616726185532763, 'mean_word_length': 2.1731225296442687, 'pattern_counts': {'":': 0, '<': 24, '<?xml version=': 0, '>': 89, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3934807', 'n_tokens_mistral': 3987, 'n_tokens_neox': 3786, 'n_words': 895}