24895189695

Committed 24 Apr 2026 02:36PM UTC coverage: 79.313% (-0.03%) from 79.342%

Build # 24895189695

Build Type

push

github

Committed by

web-flow

Commit Message

fix: retry on checkout when db_handler exits during checkout (#941)

Addresses a rare race condition.

Coverage Stats

8 of 13 new or added lines in 3 files covered. (61.54%)

1 existing line in 1 file now uncovered.

2515 of 3171 relevant lines covered (79.31%)

58036.07 hits per line

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

86.31

/lib/supavisor/db_handler.ex

defmodule Supavisor.DbHandler do
  @moduledoc """
  This module contains functions to start a connection to the database, send
  requests to the database, and handle incoming messages from clients.

  The state machine uses the Supavisor.Protocol.Server module to decode messages
  from the database and sends messages to the client socket it received on checkout.

  ## Startup modes

  DbHandler has two startup paths depending on the mode:

  ### Pool mode (transaction/session)

  Started via `start_link/1` by poolboy with a Manager config. Fetches auth
  secrets (or enters `:waiting_for_secrets` if unavailable), then connects to
  the database immediately. Workers are long-lived and shared across clients
  via checkout/checkin.

  ### Proxy mode

  Started via `start_link/1` under a DynamicSupervisor with full connection
  args. Connects to the proxy node immediately. One worker per client, with the
  DynamicSupervisor's `:max_children` enforcing the connection limit.
  """

  @behaviour :gen_statem

  require Logger
  require Supavisor
  require Supavisor.Protocol.Server, as: Server
  require Supavisor.Protocol.MessageStreamer, as: MessageStreamer

  alias Supavisor.ConnectionParameters
  alias Supavisor.Errors.CheckoutError
  alias Supavisor.Errors.CheckoutTimeoutError
  alias Supavisor.Errors.DbHandlerExitedError
  alias Supavisor.Secrets.PasswordSecrets
  alias Supavisor.Protocol.{PreparedStatements, StartupOptions}

  alias Supavisor.{
    ClientHandler,
    FeatureFlag,
    HandlerHelpers,
    Helpers,
    Monitoring.Telem,
    Protocol.BackendMessageHandler,
    Protocol.Debug,
    Protocol.MessageStreamer,
    Protocol.Server
  }

  @type state ::
          :connect
          | :connect_cooldown
          | :authentication
          | :idle
          | :busy
          | :terminating_with_error
          | :waiting_for_secrets

  @sock_closed [:tcp_closed, :ssl_closed]
  @proto [:tcp, :ssl]
  @switch_active_count Application.compile_env(:supavisor, :switch_active_count)
  @cleanup_buffer_limit 65_536
  @connect_cooldown_ms 2_500

  @auth_error_actions %{
    "28P01" => {:keep_pool, :invalidate_secrets},
    "28000" => {:keep_pool, :invalidate_secrets},
    "3D000" => {:shutdown_pool, :none},
    "42501" => {:shutdown_pool, :none},
    "22023" => {:shutdown_pool, :none},
    "57P03" => {:graceful_shutdown_pool, :none}
  }

  @doc """
  Starts a DbHandler state machine.
  """
  def start_link(config),
    do: :gen_statem.start_link(__MODULE__, config, hibernate_after: 5_000)

  def child_spec(config) do
    %{
      id: __MODULE__,
      start: {__MODULE__, :start_link, [config]},
      restart: :temporary
    }
  end

  @doc """
  Checks out a DbHandler process

  Requires a client socket. The DbHandler will forward messages directly to the
  client socket when possible.

  Returns the server socket, which the client may write messages directly to.
  """
  @spec checkout(pid(), Supavisor.sock(), pid(), Supavisor.mode(), timeout()) ::
          {:ok, Supavisor.sock()}
          | {:error, DbHandlerExitedError.t()}
          | {:error, CheckoutError.t()}
          | {:error, CheckoutTimeoutError.t()}
  def checkout(pid, sock, caller, mode, timeout \\ 15_000) do
    :gen_statem.call(pid, {:checkout, sock, caller}, timeout)
  catch
    :exit, {:timeout, _} ->
      {:error, %CheckoutTimeoutError{mode: mode, timeout_ms: timeout}}

    :exit, {reason, _} ->
      {:error, %DbHandlerExitedError{pid: pid, reason: reason}}
  end

  @doc """
  Attempts to clean up session state by sending DISCARD ALL to the database.

  The caller is responsible for ensuring that:
  - The DbHandler is NOT actively processing a query
  - The DbHandler is NOT in a transaction (no uncommitted changes)
  - The DbHandler is in session mode (not transaction mode)
  """
  @spec attempt_cleanup(pid()) :: :ok | {:error, term()}
  def attempt_cleanup(db_handler_pid) do
    :gen_statem.call(db_handler_pid, :cleanup, 5_000)
  catch
    :exit, reason ->
      {:error, {:exit, reason}}
  end

  @doc """
  Sends prepared statement packets to a DbHandler

  Different from most packets, prepared statements packets involve state at the DbHandler,
  and hence can't be sent directly to the database socket. Instead, they should be sent
  to the DbHandler through this function.
  """
  @spec handle_prepared_statement_pkts(pid, [PreparedStatements.handled_pkt()]) :: :ok
  def handle_prepared_statement_pkts(pid, pkts) do
    :gen_statem.call(pid, {:handle_ps_pkts, pkts}, 15_000)
  end

  @doc """
  Stops a DbHandler
  """
  @spec stop(pid()) :: :ok
  def stop(pid) do
    Logger.debug("DbHandler: Stop pid #{inspect(pid)}")
    :gen_statem.stop(pid, {:shutdown, :client_termination}, 5_000)
  end

  @doc """
  Notifies a DbHandler that secrets are now available
  """
  @spec notify_secrets_available(pid()) :: :ok
  def notify_secrets_available(pid) do
    :gen_statem.cast(pid, :secrets_available)
  end

  @impl true
  def init(args) do
    Process.flag(:trap_exit, true)

    {id, config} =
      case args do
        %{proxy: true} -> {args.id, args}
        %{} -> {args.id, Supavisor.Manager.get_config(args.id)}
      end

    Helpers.set_log_level(config.log_level)
    Helpers.set_max_heap_size(90)

    Logger.metadata(project: config.tenant, user: config.user, mode: config.mode)

    data = %{
      id: id,
      sock: nil,
      connection_params: config.connection_params,
      user: config.user,
      tenant: config.tenant,
      tenant_feature_flags: config.tenant_feature_flags,
      db_state: nil,
      parameter_status: %{},
      nonce: nil,
      server_proof: nil,
      stats: %{},
      prepared_statements: MapSet.new(),
      proxy: Map.get(config, :proxy, false),
      client_tls: Map.get(config, :client_tls),
      client_jit: Map.get(config, :client_jit),
      stream_state: MessageStreamer.new_stream_state(BackendMessageHandler),
      mode: config.mode,
      replica_type: config.replica_type,
      caller: nil,
      client_sock: nil,
      terminating_error: nil,
      manager_ref: nil,
      derived_secrets: nil
    }

    Telem.handler_action(:db_handler, :started, id)

    cooldown =
      if data.proxy,
        do: 0,
        else: connect_cooldown_remaining(id)

    if cooldown > 0 do
      {:ok, :connect_cooldown, data, {:state_timeout, cooldown, :connect}}
    else
      {initial_state, data, actions} = resolve_secrets(data)
      {:ok, initial_state, data, actions}
    end
  end

  @impl true
  def callback_mode, do: [:handle_event_function]

  @impl true
  def handle_event(:state_timeout, :connect, :connect_cooldown, data) do
    {next_state, data, actions} = resolve_secrets(data)
    {:next_state, next_state, data, actions}
  end

  def handle_event(:internal, :connect, :connect, %{connection_params: conn_params} = data) do
    Logger.debug("DbHandler: Try to connect to DB")

    sock_opts = [
      conn_params.ip_version,
      mode: :binary,
      packet: :raw,
      nodelay: true,
      active: false
    ]

    Telem.handler_action(:db_handler, :db_connection, data.id)

    case :gen_tcp.connect(conn_params.host, conn_params.port, sock_opts) do
      {:ok, sock} ->
        # Ensure buffer >= recbuf to avoid unnecessary copying
        # Set once at connection time as best effort; OS may adjust recbuf later via auto-tuning.
        {:ok, [{:recbuf, recbuf}]} = :inet.getopts(sock, [:recbuf])
        :ok = :inet.setopts(sock, buffer: recbuf)

        Logger.debug("DbHandler: connection_params #{inspect(conn_params, pretty: true)}")

        case try_ssl_handshake({:gen_tcp, sock}, conn_params) do
          {:ok, sock} ->
            tenant = if data.proxy, do: Supavisor.id(data.id, :tenant)

            options = %{
              "search_path" => Supavisor.id(data.id, :search_path),
              "client_tls" => if(data.proxy, do: to_string(data.client_tls)),
              "jit" => if(data.proxy, do: to_string(data.client_jit))
            }

            case send_startup(sock, conn_params, tenant, options) do
              :ok ->
                :ok = activate(sock)
                {:next_state, :authentication, %{data | sock: sock}}

              {:error, reason} ->
                Logger.error("DbHandler: Send startup error #{inspect(reason)}")
                handle_connection_failure(reason, data)
            end

          {:error, reason} ->
            Logger.error("DbHandler: Handshake error #{inspect(reason)}")
            handle_connection_failure(reason, data)
        end

      other ->
        Logger.error(
          "DbHandler: Connection failed #{inspect(other)} to #{inspect(conn_params.host)}:#{inspect(conn_params.port)}"
        )

        handle_connection_failure(other, data)
    end
  end

  def handle_event(:internal, {:terminate_with_error, error, pool_action}, _state, data) do
    Logger.debug("DbHandler: Transitioning to terminating_with_error state")

    case pool_action do
      :shutdown_pool when not data.proxy ->
        Supavisor.Manager.shutdown_with_error(data.id, error)

      :graceful_shutdown_pool when not data.proxy ->
        Supavisor.async_stop(data.id)

      _ ->
        :ok
    end

    # If not checked out yet, the postponed checkout will handle sending the error
    if data.client_sock != nil do
      encode_and_forward_error(error, data)
    end

    # Use cast to allow postponed events to be processed first
    :gen_statem.cast(self(), :finalize_termination)

    # This state will handle postponed checkout calls by returning the error
    {:next_state, :terminating_with_error, %{data | terminating_error: error}}
  end

  def handle_event(:cast, :finalize_termination, :terminating_with_error, _data) do
    Logger.debug("DbHandler: Stopping from terminating_with_error state")
    {:stop, :normal}
  end

  def handle_event(:state_timeout, :cleanup_timeout, :waiting_cleanup, _data) do
    Logger.error("DbHandler: Cleanup timeout, shutting down")
    {:stop, :normal}
  end

  def handle_event(:info, {proto, _, bin}, :authentication, data) when proto in @proto do
    {:ok, dec_pkt, _} = Server.decode(bin)
    Logger.debug("DbHandler: dec_pkt, #{inspect(dec_pkt, pretty: true)}")

    resp = Enum.reduce(dec_pkt, %{}, &handle_auth_pkts(&1, &2, data))

    case resp do
      {:authentication_sasl, nonce} ->
        {:keep_state, %{data | nonce: nonce}}

      {:authentication_server_first_message, server_proof, derived_secrets} ->
        {:keep_state,
         Map.merge(data, %{server_proof: server_proof, derived_secrets: derived_secrets})}

      %{authentication_server_final_message: _server_final} ->
        :keep_state_and_data

      %{authentication_ok: true} ->
        :keep_state_and_data

      :authentication_md5 ->
        {:keep_state, data}

      :authentication_cleartext ->
        {:keep_state, data}

      {:error_response, error} ->
        {pool_action, side_effect} = Map.get(@auth_error_actions, error["C"], {:keep_pool, :none})

        if side_effect == :invalidate_secrets do
          handle_authentication_error(data, error["M"] || "Authentication failed")
        end

        Logger.error("DbHandler: Auth error #{inspect(error)}")

        {:keep_state_and_data,
         {:next_event, :internal, {:terminate_with_error, error, pool_action}}}

      {:ready_for_query, acc} ->
        ps = acc.ps

        Logger.debug(
          "DbHandler: DB ready_for_query: #{inspect(acc.db_state)} #{inspect(ps, pretty: true)}"
        )

        if data.mode != :proxy do
          Supavisor.set_parameter_status(data.id, ps)
          cache_derived_secrets(data)
        end

        {:next_state, :idle, Map.merge(data, %{parameter_status: ps, derived_secrets: nil})}

      other ->
        Logger.error("DbHandler: Undefined auth response #{inspect(other)}")
        {:stop, :auth_error, data}
    end
  end

  # the process received message from db while idle
  def handle_event(:info, {proto, _, bin}, :idle, _data) when proto in @proto do
    Logger.debug("DbHandler: Got db response #{inspect(bin)} when idle")
    :keep_state_and_data
  end

  # forward the message to the client
  def handle_event(:info, {proto, _, bin}, :busy, %{caller: caller} = data)
      when is_pid(caller) and proto in @proto do
    Logger.debug("DbHandler: Got messages: #{Debug.packet_to_string(bin, :backend)}")

    if String.ends_with?(bin, Server.ready_for_query()) do
      ClientHandler.db_status(data.caller, :ready_for_query)
      data = handle_server_messages(bin, data)

      case data.mode do
        :transaction ->
          {_, stats} = Telem.network_usage(:db, data.sock, data.id, data.stats)
          {:next_state, :idle, %{data | stats: stats, caller: nil, client_sock: nil}}

        :proxy ->
          {:keep_state, data}

        :session ->
          {_, stats} = Telem.network_usage(:db, data.sock, data.id, data.stats)
          {:keep_state, %{data | stats: stats}}
      end
    else
      data = handle_server_messages(bin, data)
      {:keep_state, data}
    end
  end

  def handle_event(:info, {proto, _, bin}, :waiting_cleanup, %{caller: caller} = data)
      when is_pid(caller) and proto in @proto do
    buffered_bin = data.pending_bin <> bin

    cond do
      String.ends_with?(buffered_bin, Server.ready_for_query()) ->
        new_data = %{data | caller: nil, waiting_cleanup: nil, pending_bin: nil}
        {:next_state, :idle, new_data, {:reply, data.waiting_cleanup, :ok}}

      byte_size(buffered_bin) > @cleanup_buffer_limit ->
        Logger.error("DbHandler: Cleanup buffer limit exceeded, shutting down")
        {:stop, :normal}

      true ->
        {:keep_state, %{data | pending_bin: buffered_bin}}
    end
  end

  def handle_event({:call, from}, {:handle_ps_pkts, pkts}, :busy, data) do
    {iodata, data} = Enum.reduce(pkts, {[], data}, &handle_prepared_statement_pkt/2)

    {close_pkts, prepared_statements} = evict_exceeding(data)

    :ok = HandlerHelpers.sock_send(data.sock, Enum.reverse([close_pkts | iodata]))

    data = %{
      data
      | stream_state:
          Enum.reduce(close_pkts, data.stream_state, fn _, stream_state ->
            MessageStreamer.update_state(stream_state, fn queue ->
              :queue.in({:intercept, :close}, queue)
            end)
          end),
        prepared_statements: prepared_statements
    }

    {:keep_state, data, {:reply, from, :ok}}
  end

  def handle_event({:call, from}, {:checkout, _sock, _caller}, :terminating_with_error, data) do
    Logger.debug("DbHandler: checkout call during terminating_with_error, replying with error")
    error = %Supavisor.Errors.CheckoutError{pid: self(), postgres_error: data.terminating_error}
    {:keep_state_and_data, {:reply, from, {:error, error}}}
  end

  def handle_event({:call, from}, {:checkout, _sock, _caller}, :waiting_for_secrets, _data) do
    Logger.debug("DbHandler: checkout call during waiting_for_secrets, replying with error")

    postgres_error = %{
      "S" => "FATAL",
      "C" => "28P01",
      "M" =>
        "Authentication credentials are invalid. Please reconnect with fresh credentials to restore pool functionality."
    }

    error = %Supavisor.Errors.CheckoutError{pid: self(), postgres_error: postgres_error}
    {:keep_state_and_data, {:reply, from, {:error, error}}}
  end

  def handle_event({:call, from}, {:checkout, sock, caller}, state, data) do
    Logger.debug("DbHandler: checkout call when state was #{state}: #{inspect(caller)}")

    if state in [:idle, :busy] do
      Process.link(caller)

      if data.mode == :proxy do
        bin_ps = Server.encode_parameter_status(data.parameter_status)
        send(caller, {:parameter_status, bin_ps})
      end

      {:next_state, :busy, %{data | client_sock: sock, caller: caller},
       {:reply, from, {:ok, data.sock}}}
    else
      {:keep_state_and_data, :postpone}
    end
  end

  def handle_event({:call, from}, :ps, :busy, data) do
    Logger.debug("DbHandler: get parameter status")
    {:keep_state_and_data, {:reply, from, data.parameter_status}}
  end

  def handle_event({:call, from}, :cleanup, state, data) do
    Logger.debug("DbHandler: Cleanup requested, current state: #{inspect(state)}")

    cond do
      data.mode == :transaction ->
        Logger.error(
          "DbHandler: Cleanup called on transaction mode - only supported in session mode"
        )

        {:keep_state_and_data,
         {:reply, from, {:error, :cleanup_not_supported_in_transaction_mode}}}

      state in [:idle, :busy] ->
        Logger.debug("DbHandler: Starting cleanup, sending DISCARD ALL")
        msg = :pgo_protocol.encode_query_message("DISCARD ALL")
        :ok = HandlerHelpers.sock_send(data.sock, msg)

        {:next_state, :waiting_cleanup,
         Map.merge(data, %{waiting_cleanup: from, pending_bin: <<>>}),
         {:state_timeout, 5_000, :cleanup_timeout}}

      true ->
        Logger.warning("DbHandler: Cannot cleanup in state #{inspect(state)}")
        {:keep_state_and_data, {:reply, from, {:error, :cant_cleanup_now}}}
    end
  end

  def handle_event(_, {closed, _}, :busy, data) when closed in @sock_closed do
    {:stop, {:shutdown, :db_termination}, data}
  end

  def handle_event(_, {closed, _}, :authentication, data) when closed in @sock_closed do
    Logger.error("DbHandler: Db connection closed when state was authentication")

    handle_connection_failure({:error, :db_connection_closed_in_auth}, data)
  end

  def handle_event(_, {closed, _}, state, data) when closed in @sock_closed do
    if state != :terminating_with_error do
      Logger.error("DbHandler: Db connection closed when state was #{state}")
    end

    {:stop, {:shutdown, :db_termination}, data}
  end

  # linked client_handler went down
  def handle_event(_, {:EXIT, pid, reason}, _state, data) do
    if reason != :normal do
      Logger.error(
        "DbHandler: ClientHandler #{inspect(pid)} went down with reason #{inspect(reason)}"
      )
    end

    HandlerHelpers.sock_send(data.sock, Server.terminate_message())
    HandlerHelpers.sock_close(data.sock)
    {:stop, :normal}
  end

  def handle_event({:call, from}, :get_state_and_mode, state, data) do
    {:keep_state_and_data, {:reply, from, {state, data.mode}}}
  end

  def handle_event(:cast, :secrets_available, :waiting_for_secrets, data) do
    Logger.info("DbHandler: Secrets are now available, transitioning to connect state")

    Process.demonitor(data.manager_ref, [:flush])

    case get_connection_params_with_secrets(data.connection_params, data.id) do
      {:ok, conn_params_with_secrets} ->
        {:next_state, :connect,
         %{data | connection_params: conn_params_with_secrets, manager_ref: nil},
         {:next_event, :internal, :connect}}

      {:error, :no_secrets} ->
        Logger.error("DbHandler: Still no secrets available after notification")
        :keep_state_and_data
    end
  end

  def handle_event(:info, {:DOWN, ref, :process, _pid, reason}, :waiting_for_secrets, data)
      when ref == data.manager_ref do
    Logger.error("DbHandler: Manager died while waiting for secrets: #{inspect(reason)}")
    {:stop, :normal, data}
  end

  def handle_event(:info, {event, _socket}, _, data) when event in [:tcp_passive, :ssl_passive] do
    HandlerHelpers.setopts(data.sock, active: @switch_active_count)
    :keep_state_and_data
  end

  def handle_event(type, content, state, data) do
    msg = [
      {"type", type},
      {"content", content},
      {"state", state},
      {"data", data}
    ]

    Logger.debug("DbHandler: Undefined msg: #{inspect(msg, pretty: true)}")

    :keep_state_and_data
  end

  @impl true
  def terminate(_reason, :terminating_with_error, data) do
    Telem.handler_action(:db_handler, :stopped, data.id)
  end

  def terminate(reason, state, data) do
    Telem.handler_action(:db_handler, :stopped, data.id)

    case reason do
      :normal ->
        :ok

      :shutdown ->
        :ok

      reason ->
        Logger.error(
          "DbHandler: Terminating with reason #{inspect(reason)} when state was #{inspect(state)}"
        )
    end
  end

  @impl true
  def format_status(status) do
    Map.put(status, :queue, [])
  end

  @spec encode_and_forward_error(map(), map()) :: :ok | :noop
  defp encode_and_forward_error(message, data) do
    case data do
      %{client_sock: sock} when not is_nil(sock) ->
        HandlerHelpers.sock_send(
          sock,
          Server.encode_error_message(message)
        )

      _other ->
        :noop
    end
  end

  @spec try_ssl_handshake(Supavisor.tcp_sock(), ConnectionParameters.t()) ::
          {:ok, Supavisor.sock()} | {:error, term()}
  defp try_ssl_handshake(sock, %ConnectionParameters{upstream_ssl: true} = conn_params) do
    case HandlerHelpers.sock_send(sock, Server.ssl_request()) do
      :ok -> ssl_recv(sock, conn_params)
      error -> error
    end
  end

  defp try_ssl_handshake(sock, _), do: {:ok, sock}

  @spec ssl_recv(Supavisor.tcp_sock(), ConnectionParameters.t()) ::
          {:ok, Supavisor.ssl_sock()} | {:error, term}
  defp ssl_recv({:gen_tcp, sock} = s, conn_params) do
    case :gen_tcp.recv(sock, 1, 15_000) do
      {:ok, <<?S>>} -> ssl_connect(s, conn_params)
      {:ok, <<?N>>} -> {:error, :ssl_not_available}
      {:error, _} = error -> error
    end
  end

  @spec ssl_connect(Supavisor.tcp_sock(), ConnectionParameters.t(), pos_integer) ::
          {:ok, Supavisor.ssl_sock()} | {:error, term}
  defp ssl_connect({:gen_tcp, sock}, conn_params, timeout \\ 5000) do
    opts =
      case conn_params.upstream_verify do
        :peer ->
          [
            verify: :verify_peer,
            cacerts: [conn_params.upstream_tls_ca],
            # unclear behavior on pg14
            server_name_indication: conn_params.sni_hostname || conn_params.host,
            customize_hostname_check: [{:match_fun, fn _, _ -> true end}],
            receiver_spawn_opts: [min_heap_size: 2048]
          ]

        :none ->
          [verify: :verify_none, receiver_spawn_opts: [min_heap_size: 2048]]
      end

    case :ssl.connect(sock, opts, timeout) do
      {:ok, ssl_sock} ->
        {:ok, {:ssl, ssl_sock}}

      {:error, reason} ->
        {:error, reason}
    end
  end

  @spec send_startup(Supavisor.sock(), ConnectionParameters.t(), String.t() | nil, map()) ::
          :ok | {:error, term}
  def send_startup(sock, conn_params, tenant, options) do
    user =
      if is_nil(tenant),
        do: conn_params.secrets.user,
        else: "#{conn_params.secrets.user}.#{tenant}"

    options = Map.reject(options, fn {_k, v} -> is_nil(v) end)

    msg =
      :pgo_protocol.encode_startup_message(
        [
          {"user", user},
          {"database", conn_params.database},
          {"application_name", conn_params.application_name}
        ] ++
          if(options != %{},
            do: [{"options", StartupOptions.encode(options)}],
            else: []
          )
      )

    HandlerHelpers.sock_send(sock, msg)
  end

  @spec activate(Supavisor.sock()) :: :ok | {:error, term}
  defp activate({:gen_tcp, sock}) do
    :inet.setopts(sock, active: @switch_active_count)
  end

  defp activate({:ssl, sock}) do
    :ssl.setopts(sock, active: @switch_active_count)
  end

  @spec handle_auth_pkts(map(), map(), map()) :: any()
  defp handle_auth_pkts(%{tag: :parameter_status, payload: {k, v}}, acc, _),
    do: update_in(acc, [:ps], fn ps -> Map.put(ps || %{}, k, v) end)

  defp handle_auth_pkts(%{tag: :ready_for_query, payload: db_state}, acc, _),
    do: {:ready_for_query, Map.put(acc, :db_state, db_state)}

  defp handle_auth_pkts(%{tag: :backend_key_data, payload: payload}, acc, data) do
    if data.mode != :proxy do
      Logger.metadata(backend_pid: payload[:pid])
    end

    key = self()

    conn = %{
      host: data.connection_params.host,
      port: data.connection_params.port,
      ip_version: data.connection_params.ip_version
    }

    Registry.register(Supavisor.Registry.PoolPids, key, Map.merge(payload, conn))
    Logger.debug("DbHandler: Backend #{inspect(key)} data: #{inspect(payload)}")
    Map.put(acc, :backend_key_data, payload)
  end

  defp handle_auth_pkts(%{payload: {:authentication_sasl_password, methods_b}}, _, data) do
    nonce =
      case Server.decode_string(methods_b) do
        {:ok, req_method, _} ->
          Logger.debug("DbHandler: SASL method #{inspect(req_method)}")
          nonce = :pgo_scram.get_nonce(16)
          user = data.connection_params.secrets.user
          client_first = :pgo_scram.get_client_first(user, nonce)
          client_first_size = IO.iodata_length(client_first)

          sasl_initial_response = [
            "SCRAM-SHA-256",
            0,
            <<client_first_size::32-integer>>,
            client_first
          ]

          bin = :pgo_protocol.encode_scram_response_message(sasl_initial_response)
          :ok = HandlerHelpers.sock_send(data.sock, bin)
          nonce

        other ->
          Logger.error("DbHandler: Undefined sasl method #{inspect(other)}")
          nil
      end

    {:authentication_sasl, nonce}
  end

  defp handle_auth_pkts(
         %{payload: {:authentication_server_first_message, server_first}},
         _,
         data
       ) do
    nonce = data.nonce
    server_first_parts = Helpers.parse_server_first(server_first, nonce)

    secrets = data.connection_params.secrets

    {client_final_message, server_proof, derived_secrets} =
      Helpers.get_client_final(
        secrets,
        server_first_parts,
        nonce,
        secrets.user,
        "biws"
      )

    bin = :pgo_protocol.encode_scram_response_message(client_final_message)
    :ok = HandlerHelpers.sock_send(data.sock, bin)

    {:authentication_server_first_message, server_proof, derived_secrets}
  end

  defp handle_auth_pkts(
         %{payload: {:authentication_server_final_message, server_final}},
         acc,
         _data
       ),
       do: Map.put(acc, :authentication_server_final_message, server_final)

  defp handle_auth_pkts(
         %{payload: :authentication_ok},
         acc,
         _data
       ),
       do: Map.put(acc, :authentication_ok, true)

  defp handle_auth_pkts(%{payload: {:authentication_md5_password, salt}} = dec_pkt, _, data) do
    Logger.debug("DbHandler: dec_pkt, #{inspect(dec_pkt, pretty: true)}")

    %PasswordSecrets{password: password, user: user} = data.connection_params.secrets

    digest = Helpers.md5([password, user])

    payload = ["md5", Helpers.md5([digest, salt]), 0]
    bin = [?p, <<IO.iodata_length(payload) + 4::signed-32>>, payload]
    :ok = HandlerHelpers.sock_send(data.sock, bin)
    :authentication_md5
  end

  defp handle_auth_pkts(%{payload: :authentication_cleartext_password} = dec_pkt, _, data) do
    Logger.debug("DbHandler: dec_pkt, #{inspect(dec_pkt, pretty: true)}")

    %PasswordSecrets{password: password} = data.connection_params.secrets
    payload = <<password::binary, 0>>
    bin = [?p, <<IO.iodata_length(payload) + 4::signed-32>>, payload]
    :ok = HandlerHelpers.sock_send(data.sock, bin)
    :authentication_cleartext
  end

  defp handle_auth_pkts(%{tag: :error_response, payload: error}, _acc, _data),
    do: {:error_response, error}

  defp handle_auth_pkts(_e, acc, _data), do: acc

  defp cache_derived_secrets(%{id: id, derived_secrets: derived_secrets})
       when not is_nil(derived_secrets) do
    Supavisor.UpstreamAuthentication.put_upstream_auth_secrets(id, derived_secrets)
  end

  defp cache_derived_secrets(_data), do: :ok

  @spec handle_authentication_error(map(), String.t()) :: any()
  defp handle_authentication_error(%{mode: :proxy}, _reason), do: :ok

  defp handle_authentication_error(%{mode: _other} = data, _reason) do
    tenant = Supavisor.id(data.id, :tenant)
    Supavisor.ClientAuthentication.invalidate_global(tenant, data.user)
    Supavisor.UpstreamAuthentication.delete_upstream_auth_secrets(data.id)
  end

  @spec handle_server_messages(binary(), map()) :: map()
  defp handle_server_messages(bin, data) do
    if FeatureFlag.enabled?(data.tenant_feature_flags, "named_prepared_statements") do
      {:ok, updated_data, packets_to_send} = process_backend_streaming(bin, data)

      if packets_to_send != [] do
        HandlerHelpers.sock_send(data.client_sock, packets_to_send)
      end

      updated_data
    else
      HandlerHelpers.sock_send(data.client_sock, bin)

      data
    end
  end

  # If the prepared statement exists for us, it exists for the server, so we just send the
  # bind to the socket. If it doesn't, we must send the parse pkt first.
  #
  # If we received a bind without a parse, we need to intercept the parse response, otherwise,
  # the client will receive an unexpected message.
  defp handle_prepared_statement_pkt({:bind_pkt, stmt_name, pkt, parse_pkt}, {iodata, data}) do
    if stmt_name in data.prepared_statements do
      {[pkt | iodata], data}
    else
      new_data = %{
        data
        | stream_state:
            MessageStreamer.update_state(data.stream_state, fn queue ->
              :queue.in({:intercept, :parse}, queue)
            end),
          prepared_statements: MapSet.put(data.prepared_statements, stmt_name)
      }

      {[[parse_pkt, pkt] | iodata], new_data}
    end
  end

  defp handle_prepared_statement_pkt({:close_pkt, stmt_name, pkt}, {iodata, data}) do
    {[pkt | iodata],
     %{
       data
       | prepared_statements: MapSet.delete(data.prepared_statements, stmt_name),
         stream_state:
           MessageStreamer.update_state(data.stream_state, fn queue ->
             :queue.in({:forward, :close}, queue)
           end)
     }}
  end

  defp handle_prepared_statement_pkt({:describe_pkt, _stmt_name, pkt}, {iodata, data}) do
    {[pkt | iodata], data}
  end

  # If we stop generating unique id per statement, and instead do deterministic ids,
  # we need to potentially drop parse pkts and return a parse response
  defp handle_prepared_statement_pkt({:parse_pkt, stmt_name, pkt}, {iodata, data}) do
    if stmt_name in data.prepared_statements do
      {iodata,
       %{
         data
         | stream_state:
             MessageStreamer.update_state(data.stream_state, fn queue ->
               :queue.in({:inject, :parse}, queue)
             end)
       }}
    else
      prepared_statements = MapSet.put(data.prepared_statements, stmt_name)

      {[pkt | iodata],
       %{
         data
         | prepared_statements: prepared_statements,
           stream_state:
             MessageStreamer.update_state(data.stream_state, fn queue ->
               :queue.in({:forward, :parse}, queue)
             end)
       }}
    end
  end

  defp evict_exceeding(%{prepared_statements: prepared_statements, id: id}) do
    limit = PreparedStatements.backend_limit()

    if MapSet.size(prepared_statements) >= limit do
      count = div(limit, 5)
      to_remove = Enum.take_random(prepared_statements, count) |> MapSet.new()
      close_pkts = Enum.map(to_remove, &PreparedStatements.build_close_pkt/1)
      prepared_statements = MapSet.difference(prepared_statements, to_remove)
      Telem.prepared_statements_evicted(count, id)

      {close_pkts, prepared_statements}
    else
      {[], prepared_statements}
    end
  end

  defp process_backend_streaming(bin, data) do
    case MessageStreamer.handle_packets(data.stream_state, bin) do
      {:ok, new_stream_state, packets} ->
        updated_data = %{data | stream_state: new_stream_state}
        {:ok, updated_data, packets}

      err ->
        err
    end
  end

  defp get_connection_params_with_secrets(conn_params, id) do
    case Supavisor.UpstreamAuthentication.get_upstream_auth_secrets(id) do
      {:ok, upstream_auth_secrets} ->
        {:ok, %{conn_params | secrets: upstream_auth_secrets}}

      _other ->
        {:error, :no_secrets}
    end
  end

  defp handle_connection_failure(reason, data) do
    if not data.proxy do
      Supavisor.CircuitBreaker.record_failure(data.tenant, :db_connection)
      Supavisor.TenantCache.put_last_connect_failure(data.id, System.monotonic_time(:millisecond))
    end

    error = %{
      "S" => "FATAL",
      "C" => "08006",
      "M" => "Failed to connect to database: #{inspect(reason)}"
    }

    {:keep_state_and_data, {:next_event, :internal, {:terminate_with_error, error, :keep_pool}}}
  end

  defp resolve_secrets(data) do
    if data.proxy do
      {:connect, data, {:next_event, :internal, :connect}}
    else
      case get_connection_params_with_secrets(data.connection_params, data.id) do
        {:ok, conn_params_with_secrets} ->
          {:connect, %{data | connection_params: conn_params_with_secrets},
           {:next_event, :internal, :connect}}

        {:error, :no_secrets} ->
          Logger.warning("DbHandler: Secrets not available, entering waiting state")
          manager_pid = Supavisor.get_local_manager(data.id)
          ref = Process.monitor(manager_pid)
          Supavisor.Manager.register_waiting_for_secrets(data.id, self())
          {:waiting_for_secrets, %{data | manager_ref: ref}, []}
      end
    end
  end

  defp connect_cooldown_remaining(id) do
    case Supavisor.TenantCache.get_last_connect_failure(id) do
      nil ->
        0

      last_failure ->
        elapsed = System.monotonic_time(:millisecond) - last_failure
        @connect_cooldown_ms - elapsed
    end
  end
end

supabase / supavisor / 24895189695

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