diff --git a/lib/supavisor/client_handler.ex b/lib/supavisor/client_handler.ex index ec95c413e..dfb26ce9b 100644 --- a/lib/supavisor/client_handler.ex +++ b/lib/supavisor/client_handler.ex @@ -45,7 +45,8 @@ defmodule Supavisor.ClientHandler do Data, Error, ProtocolHelpers, - Proxy + Proxy, + StartupParams } alias Supavisor.Protocol.{FrontendMessageHandler, MessageStreamer, StartupOptions} @@ -216,16 +217,8 @@ defmodule Supavisor.ClientHandler do def handle_event(:info, {_, _, bin}, :handshake, data) do case ProtocolHelpers.parse_startup_packet(bin) do - {:ok, {type, {user, tenant_or_alias, db_name, search_path, jit, client_tls, client_ip}}, - app_name, log_level, invalid} -> - event = - {:hello, - {type, {user, tenant_or_alias, db_name, search_path, jit, client_tls, client_ip}}} - - if log_level, do: Logger.put_process_level(self(), log_level) - - {:keep_state, %{data | app_name: app_name, invalid_options: invalid}, - {:next_event, :internal, event}} + {:ok, %StartupParams{} = startup_params} -> + {:keep_state_and_data, {:next_event, :internal, {:hello, startup_params}}} {:error, exception} -> Error.terminate_with_error(data, exception, :handshake) @@ -234,12 +227,24 @@ defmodule Supavisor.ClientHandler do def handle_event( :internal, - {:hello, - {type, {user, tenant_or_alias, db_name, search_path, client_jit, client_tls, client_ip}} = - hello_args}, + {:hello, %StartupParams{} = startup_params}, :handshake, %{sock: sock} = data ) do + %StartupParams{ + type: type, + user: user, + tenant_or_alias: tenant_or_alias, + db_name: db_name, + search_path: search_path, + jit: client_jit, + client_tls: client_tls, + client_ip: client_ip, + app_name: app_name, + log_level: log_level, + invalid_options: invalid_options + } = startup_params + sni_hostname = HandlerHelpers.try_get_sni(sock) # When receiving a proxied connection on a local listener, client_tls and @@ -247,14 +252,16 @@ defmodule Supavisor.ClientHandler do # peer is the forwarding node). Otherwise, use what we observed on the socket. effective_ssl = if(data.local && client_tls, do: client_tls, else: data.ssl) peer_ip = ProtocolHelpers.effective_peer_ip(data.local, client_ip, data.peer_ip) - data = %{data | peer_ip: peer_ip} + data = %{data | peer_ip: peer_ip, app_name: app_name, invalid_options: invalid_options} + + if log_level, do: Logger.put_process_level(self(), log_level) Logger.metadata( project: tenant_or_alias, user: user, mode: data.mode, type: type, - app_name: data.app_name, + app_name: app_name, db_name: db_name, peer_ip: peer_ip, tls: effective_ssl @@ -295,7 +302,7 @@ defmodule Supavisor.ClientHandler do {:next_event, :internal, {:start_authentication, auth_method, info}}} else {:error, %MaxConnectionsError{} = exception} -> - wait_for_slot_or_terminate(%{data | id: id}, {:hello, hello_args}, exception) + wait_for_slot_or_terminate(%{data | id: id}, {:hello, startup_params}, exception) {:error, exception} when is_exception(exception) -> Error.terminate_with_error(%{data | id: id}, exception, :handshake) diff --git a/lib/supavisor/client_handler/protocol_helpers.ex b/lib/supavisor/client_handler/protocol_helpers.ex index bc9ea8da1..b417aea81 100644 --- a/lib/supavisor/client_handler/protocol_helpers.ex +++ b/lib/supavisor/client_handler/protocol_helpers.ex @@ -14,6 +14,7 @@ defmodule Supavisor.ClientHandler.ProtocolHelpers do require Logger alias Supavisor.{ + ClientHandler.StartupParams, Errors.InvalidUserInfoError, Errors.StartupMessageError, Errors.MaxPreparedStatementsError, @@ -37,11 +38,6 @@ defmodule Supavisor.ClientHandler.ProtocolHelpers do | {:error, DuplicatePreparedStatementError.t()} | {:error, PreparedStatementNotFoundError.t()} - @type startup_message_data() :: - {atom(), - {String.t(), String.t(), String.t() | nil, String.t() | nil, boolean(), - boolean() | nil, String.t() | nil}} - ## Startup Packet Processing @doc """ @@ -50,18 +46,21 @@ defmodule Supavisor.ClientHandler.ProtocolHelpers do Returns parsed user info, application name, log level, and list of invalid options. """ @spec parse_startup_packet(binary()) :: - {:ok, startup_message_data(), String.t() | nil, Logger.level() | nil, - [{String.t(), String.t()}]} + {:ok, StartupParams.t()} | {:error, StartupMessageError.t() | InvalidUserInfoError.t()} def parse_startup_packet(bin) do with {:ok, hello} <- Client.decode_startup_packet(bin), {options, invalid} = StartupOptions.validate(hello.payload["options"] || %{}), - {:ok, user_info} <- extract_and_validate_user_info(hello.payload, options) do + {:ok, params} <- extract_and_validate_user_info(hello.payload, options) do Logger.debug("ClientHandler: Client startup message: #{inspect(hello)}") - app_name = normalize_app_name(hello.payload["application_name"]) - log_level = options["log_level"] - {:ok, user_info, app_name, log_level, invalid} + {:ok, + %{ + params + | app_name: normalize_app_name(hello.payload["application_name"]), + log_level: options["log_level"], + invalid_options: invalid + }} end end @@ -69,19 +68,25 @@ defmodule Supavisor.ClientHandler.ProtocolHelpers do Extracts and validates user information from startup payload. """ @spec extract_and_validate_user_info(map(), map()) :: - {:ok, startup_message_data()} + {:ok, StartupParams.t()} | {:error, InvalidUserInfoError.t()} def extract_and_validate_user_info(payload, options) do {type, {user, tenant_or_alias, db_name}} = HandlerHelpers.parse_user_info(payload) if Helpers.validate_name(user) and (is_nil(db_name) or Helpers.validate_name(db_name)) do - search_path = payload["search_path"] || options["search_path"] - jit = Map.get(options, "jit", false) - client_tls = Map.get(options, "client_tls") - # Set by a peer node when it forwards a proxied connection; carries the - # original client's IP. Only honored on local listeners, see effective_peer_ip/3. - client_ip = options["client_ip"] - {:ok, {type, {user, tenant_or_alias, db_name, search_path, jit, client_tls, client_ip}}} + {:ok, + %StartupParams{ + type: type, + user: user, + tenant_or_alias: tenant_or_alias, + db_name: db_name, + search_path: payload["search_path"] || options["search_path"], + jit: Map.get(options, "jit", false), + client_tls: Map.get(options, "client_tls"), + # Set by a peer node when it forwards a proxied connection; carries the + # original client's IP. Only honored on local listeners, see effective_peer_ip/3. + client_ip: options["client_ip"] + }} else {:error, %InvalidUserInfoError{user: user, db_name: db_name}} end diff --git a/lib/supavisor/client_handler/startup_params.ex b/lib/supavisor/client_handler/startup_params.ex new file mode 100644 index 000000000..7e59f9f0c --- /dev/null +++ b/lib/supavisor/client_handler/startup_params.ex @@ -0,0 +1,34 @@ +defmodule Supavisor.ClientHandler.StartupParams do + @moduledoc """ + Connection parameters parsed from a client's startup packet. + """ + + @type t :: %__MODULE__{ + type: :single | :cluster, + user: String.t(), + tenant_or_alias: String.t() | nil, + db_name: String.t() | nil, + search_path: String.t() | nil, + jit: boolean(), + client_tls: boolean() | nil, + client_ip: String.t() | nil, + app_name: String.t() | nil, + log_level: Logger.level() | nil, + invalid_options: [{String.t(), String.t()}] + } + + @enforce_keys [:type, :user] + defstruct [ + :type, + :user, + :tenant_or_alias, + :db_name, + :search_path, + :client_tls, + :client_ip, + :app_name, + :log_level, + jit: false, + invalid_options: [] + ] +end diff --git a/test/supavisor/client_handler/protocol_helpers_test.exs b/test/supavisor/client_handler/protocol_helpers_test.exs index e6b1ed2e3..84817eab9 100644 --- a/test/supavisor/client_handler/protocol_helpers_test.exs +++ b/test/supavisor/client_handler/protocol_helpers_test.exs @@ -1,23 +1,32 @@ defmodule Supavisor.ClientHandler.ProtocolHelpersTest do use ExUnit.Case, async: true - alias Supavisor.ClientHandler.ProtocolHelpers + alias Supavisor.ClientHandler.{ProtocolHelpers, StartupParams} alias Supavisor.Protocol.StartupOptions describe "parse_startup_packet/1" do test "drops an invalid option and reports it" do bin = startup_packet([{"user", "postgres.some_tenant"}, {"options", "-c jit=maybe"}]) - assert {:ok, {_type, {"postgres", "some_tenant", _db, _sp, false, _tls, _ip}}, _app, _log, - [{"jit", "maybe"}]} = ProtocolHelpers.parse_startup_packet(bin) + assert {:ok, + %StartupParams{ + user: "postgres", + tenant_or_alias: "some_tenant", + jit: false, + invalid_options: [{"jit", "maybe"}] + }} = ProtocolHelpers.parse_startup_packet(bin) end test "accepts and type-converts valid options" do bin = startup_packet([{"user", "postgres.some_tenant"}, {"options", "-c jit=1"}]) - assert {:ok, {_type, {"postgres", "some_tenant", _db, _sp, true, _tls, _ip}}, _app, _log, - []} = - ProtocolHelpers.parse_startup_packet(bin) + assert {:ok, + %StartupParams{ + user: "postgres", + tenant_or_alias: "some_tenant", + jit: true, + invalid_options: [] + }} = ProtocolHelpers.parse_startup_packet(bin) end end @@ -25,23 +34,47 @@ defmodule Supavisor.ClientHandler.ProtocolHelpersTest do test "returns nil client_ip when the option is absent" do payload = %{"user" => "postgres.some_tenant"} - assert {:ok, {_type, {"postgres", "some_tenant", nil, nil, true, nil, nil}}} = - ProtocolHelpers.extract_and_validate_user_info(payload, %{"jit" => true}) + assert {:ok, + %StartupParams{ + user: "postgres", + tenant_or_alias: "some_tenant", + db_name: nil, + search_path: nil, + jit: true, + client_tls: nil, + client_ip: nil + }} = ProtocolHelpers.extract_and_validate_user_info(payload, %{"jit" => true}) end test "returns nil client_ip when there are no options" do payload = %{"user" => "postgres.some_tenant"} - assert {:ok, {_type, {"postgres", "some_tenant", nil, nil, false, nil, nil}}} = - ProtocolHelpers.extract_and_validate_user_info(payload, %{}) + assert {:ok, + %StartupParams{ + user: "postgres", + tenant_or_alias: "some_tenant", + db_name: nil, + search_path: nil, + jit: false, + client_tls: nil, + client_ip: nil + }} = ProtocolHelpers.extract_and_validate_user_info(payload, %{}) end test "extracts client_ip alongside jit and client_tls" do payload = %{"user" => "postgres.some_tenant"} options = %{"jit" => true, "client_tls" => true, "client_ip" => "203.0.113.9"} - assert {:ok, {_type, {"postgres", "some_tenant", nil, nil, true, true, "203.0.113.9"}}} = - ProtocolHelpers.extract_and_validate_user_info(payload, options) + assert {:ok, + %StartupParams{ + user: "postgres", + tenant_or_alias: "some_tenant", + db_name: nil, + search_path: nil, + jit: true, + client_tls: true, + client_ip: "203.0.113.9" + }} = ProtocolHelpers.extract_and_validate_user_info(payload, options) end test "round-trips client_ip through the startup options wire format" do @@ -57,8 +90,16 @@ defmodule Supavisor.ClientHandler.ProtocolHelpersTest do {options, []} = StartupOptions.validate(StartupOptions.parse(encoded)) payload = %{"user" => "postgres.some_tenant"} - assert {:ok, {_type, {"postgres", "some_tenant", nil, nil, true, true, "2001:db8::1"}}} = - ProtocolHelpers.extract_and_validate_user_info(payload, options) + assert {:ok, + %StartupParams{ + user: "postgres", + tenant_or_alias: "some_tenant", + db_name: nil, + search_path: nil, + jit: true, + client_tls: true, + client_ip: "2001:db8::1" + }} = ProtocolHelpers.extract_and_validate_user_info(payload, options) end end diff --git a/test/supavisor/client_handler_test.exs b/test/supavisor/client_handler_test.exs index fab70c580..606ccea23 100644 --- a/test/supavisor/client_handler_test.exs +++ b/test/supavisor/client_handler_test.exs @@ -3,6 +3,7 @@ defmodule Supavisor.ClientHandlerTest do alias Supavisor.Protocol.FrontendMessageHandler alias Supavisor.Protocol.MessageStreamer + alias Supavisor.ClientHandler.StartupParams @subject Supavisor.ClientHandler @@ -76,7 +77,15 @@ defmodule Supavisor.ClientHandlerTest do setup do %{ exception: MaxConnectionsError.new(:transaction, 2), - retry_event: {:hello, {:single, {"user", "tenant", "postgres", nil, false, false, nil}}}, + retry_event: + {:hello, + %StartupParams{ + type: :single, + user: "user", + tenant_or_alias: "tenant", + db_name: "postgres", + client_tls: false + }}, budget: Application.get_env(:supavisor, :admission_retries) } end @@ -213,17 +222,44 @@ defmodule Supavisor.ClientHandlerTest do describe "startup packet log_level option" do test "sets process log level from options" do - bin = - <<79::32, - "\x00\x03\x00\x00user\x00postgres.dev_tenant\x00database\x00postgres\x00options\x00-c log_level=debug\x00\x00">> + # No tenant in the user name, so the hello handler fails before any DB lookup. + payload = + "\x00\x03\x00\x00user\x00log_level_user\x00database\x00postgres\x00options\x00-c log_level=debug\x00\x00" - data = %{sock: {:gen_tcp, :fake_port}, id: "test", app_name: nil, invalid_options: []} + bin = <> - assert {:keep_state, %{app_name: ""}, - {:next_event, :internal, - {:hello, {:single, {"postgres", "dev_tenant", "postgres", nil, false, nil, nil}}}}} = - @subject.handle_event(:info, {:tcp, :fake_port, bin}, :handshake, data) + {sock, _recv} = sockpair() + + data = %{ + sock: {:gen_tcp, sock}, + id: nil, + local: false, + ssl: false, + peer_ip: "127.0.0.1", + mode: :transaction, + app_name: nil, + invalid_options: [] + } + assert {:keep_state_and_data, + {:next_event, :internal, + {:hello, + %StartupParams{ + type: :single, + user: "log_level_user", + tenant_or_alias: nil, + db_name: "postgres", + search_path: nil, + jit: false, + client_tls: nil, + client_ip: nil, + app_name: "", + log_level: :debug, + invalid_options: [] + }} = hello}} = + @subject.handle_event(:info, {:tcp, sock, bin}, :handshake, data) + + assert {:stop, :normal} = @subject.handle_event(:internal, hello, :handshake, data) assert Logger.get_process_level(self()) == :debug end end