feat: use simple_connection for auth_query

VERSION

@@ -1 +1 @@
-1.1.53
+1.1.54

lib/single_connection.ex

@@ -0,0 +1,49 @@
+defmodule Supavisor.SingleConnection do
+  require Logger
+  @behaviour Postgrex.SimpleConnection
+
+  def connect(conf), do: Postgrex.SimpleConnection.start_link(__MODULE__, conf, conf)
+
+  @impl true
+  def init(args) do
+    Logger.debug("init args: #{inspect(args, pretty: true)}")
+    Process.monitor(args[:caller])
+    # put the hostname in the process dictionary to be able to find it in an emergency
+    Process.put(:auth_host, args[:hostname])
+    {:ok, %{from: nil, caller: args[:caller]}}
+  end
+
+  @impl true
+  def handle_call({:query, query}, from, state), do: {:query, query, %{state | from: from}}
+
+  def handle_result(results, state) when is_list(results) do
+    result =
+      case results do
+        [%Postgrex.Result{} = res] -> res
+        other -> other
+      end
+
+    Postgrex.SimpleConnection.reply(state.from, result)
+    {:noreply, state}
+  end
+
+  @impl true
+  def handle_result(%Postgrex.Error{} = error, state) do
+    Postgrex.SimpleConnection.reply(state.from, error)
+    {:noreply, state}
+  end
+
+  @impl true
+  def handle_info({:DOWN, _, _, caller, _}, %{caller: caller} = state) do
+    Logger.notice("Caller #{inspect(caller)} is down")
+    {:stop, state}
+  end
+
+  def handle_info(msg, state) do
+    Logger.error("Undefined message #{inspect(msg, pretty: true)}")
+    {:noreply, state}
+  end
+
+  @impl true
+  def notify(_, _, _), do: :ok
+end

lib/supavisor/client_handler.ex

@@ -15,7 +15,7 @@ defmodule Supavisor.ClientHandler do
   alias Supavisor.DbHandler, as: Db
   alias Supavisor.Helpers, as: H
   alias Supavisor.HandlerHelpers, as: HH
-  alias Supavisor.{Tenants, Monitoring.Telem, Protocol.Client, Protocol.Server}
+  alias Supavisor.{Tenants, Monitoring.Telem, Protocol.Client, Protocol.Server, SingleConnection}
 
   @impl true
   def start_link(ref, _sock, transport, opts) do
@@ -822,6 +822,8 @@ defmodule Supavisor.ClientHandler do
 
   @spec get_secrets(map, String.t()) :: {:ok, {:auth_query, fun()}} | {:error, term()}
   def get_secrets(%{user: user, tenant: tenant}, db_user) do
+    Logger.info("ClientHandler: Get secrets started")
+
     ssl_opts =
       if tenant.upstream_ssl and tenant.upstream_verify == "peer" do
         [
@@ -833,7 +835,7 @@ defmodule Supavisor.ClientHandler do
       end
 
     {:ok, conn} =
-      Postgrex.start_link(
+      SingleConnection.connect(
         hostname: tenant.db_host,
         port: tenant.db_port,
         database: tenant.db_database,
@@ -846,9 +848,14 @@ defmodule Supavisor.ClientHandler do
         ],
         queue_target: 1_000,
         queue_interval: 5_000,
-        ssl_opts: ssl_opts || []
+        ssl_opts: ssl_opts || [],
+        caller: self()
       )
 
+    Logger.debug(
+      "ClientHandler: Connected to db #{tenant.db_host} #{tenant.db_port} #{tenant.db_database} #{user.db_user}"
+    )
+
     resp =
       case H.get_user_secret(conn, tenant.auth_query, db_user) do
         {:ok, secret} ->
@@ -859,7 +866,7 @@ defmodule Supavisor.ClientHandler do
           {:error, reason}
       end
 
-    GenServer.stop(conn, :normal)
+    Logger.info("ClientHandler: Get secrets finished")
     resp
   end
 

lib/supavisor/helpers.ex

@@ -88,7 +88,11 @@ defmodule Supavisor.Helpers do
   @spec get_user_secret(pid(), String.t(), String.t()) :: {:ok, map()} | {:error, String.t()}
   def get_user_secret(conn, auth_query, user) do
     try do
-      Postgrex.query!(conn, auth_query, [user])
+      # sanitize the user input by removing all characters that are not alphanumeric or underscores
+      user = String.replace(user, ~r/[^a-zA-Z0-9_]/, "")
+      auth_query = String.replace(auth_query, "$1", "'#{user}'")
+
+      Postgrex.SimpleConnection.call(conn, {:query, auth_query})
     catch
       _error, reason ->
         {:error, "Authentication query failed: #{inspect(reason)}"}

test/integration/single_connection_test.exs

@@ -0,0 +1,36 @@
+defmodule Supavisor.Integration.SingleConnectionTest do
+  require Logger
+  use Supavisor.DataCase, async: true
+  alias Postgrex, as: P
+
+  @tenant "proxy_tenant1"
+
+  test "connects to database and executes a simple query" do
+    db_conf = Application.get_env(:supavisor, Repo)
+
+    args = [
+      hostname: db_conf[:hostname],
+      port: Application.get_env(:supavisor, :proxy_port_transaction),
+      database: "postgres",
+      password: db_conf[:password],
+      username: "transaction.#{@tenant}"
+    ]
+
+    spawn(fn ->
+      {:ok, pid} =
+        args
+        |> Keyword.put_new(:caller, self())
+        |> Supavisor.SingleConnection.connect()
+
+      assert %Postgrex.Result{rows: [["1"]]} =
+               Postgrex.SimpleConnection.call(pid, {:query, "SELECT 1"})
+    end)
+
+    :timer.sleep(250)
+
+    # check that the connection dies after the caller dies
+    assert Enum.filter(Process.list(), fn pid ->
+             Process.info(pid)[:dictionary][:auth_host] == db_conf[:hostname]
+           end) == []
+  end
+end