From 6c6129fbb7b63b9f26a730e2913caa5bcefbdee0 Mon Sep 17 00:00:00 2001
From: Stas <sts@abc3.dev>
Date: Mon, 27 May 2024 17:48:35 +0200
Subject: [PATCH 1/3] feat: use simple_connection for auth_query

---
 VERSION                                     |  2 +-
 lib/single_connection.ex                    | 49 +++++++++++++++++++++
 lib/supavisor/client_handler.ex             | 12 +++--
 lib/supavisor/helpers.ex                    |  4 +-
 test/integration/single_connection_test.exs | 36 +++++++++++++++
 5 files changed, 97 insertions(+), 6 deletions(-)
 create mode 100644 lib/single_connection.ex
 create mode 100644 test/integration/single_connection_test.exs

diff --git a/VERSION b/VERSION
index 36a0393f..baeb5d0c 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-1.1.53
+1.1.54
diff --git a/lib/single_connection.ex b/lib/single_connection.ex
new file mode 100644
index 00000000..ea2fac35
--- /dev/null
+++ b/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
diff --git a/lib/supavisor/client_handler.ex b/lib/supavisor/client_handler.ex
index e1076d39..31546d68 100644
--- a/lib/supavisor/client_handler.ex
+++ b/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
@@ -833,7 +833,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 +846,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 +864,6 @@ defmodule Supavisor.ClientHandler do
           {:error, reason}
       end
 
-    GenServer.stop(conn, :normal)
     resp
   end
 
diff --git a/lib/supavisor/helpers.ex b/lib/supavisor/helpers.ex
index 26b97ef4..0a2d3752 100644
--- a/lib/supavisor/helpers.ex
+++ b/lib/supavisor/helpers.ex
@@ -88,7 +88,9 @@ 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])
+      # Postgrex.query!(conn, auth_query, [user])
+      auth_query = String.replace(auth_query, "$1", "'#{user}'")
+      Postgrex.SimpleConnection.call(conn, {:query, auth_query})
     catch
       _error, reason ->
         {:error, "Authentication query failed: #{inspect(reason)}"}
diff --git a/test/integration/single_connection_test.exs b/test/integration/single_connection_test.exs
new file mode 100644
index 00000000..3b5928fb
--- /dev/null
+++ b/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

From fbde9dae2c386aa5fcd7a42805d47d7e552ac6b3 Mon Sep 17 00:00:00 2001
From: Stas <sts@abc3.dev>
Date: Tue, 28 May 2024 14:27:14 +0200
Subject: [PATCH 2/3] sanitize the user input

---
 lib/supavisor/helpers.ex | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/lib/supavisor/helpers.ex b/lib/supavisor/helpers.ex
index 0a2d3752..7b26125e 100644
--- a/lib/supavisor/helpers.ex
+++ b/lib/supavisor/helpers.ex
@@ -88,8 +88,10 @@ 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 ->

From 1e4e9e7061b8898fd597414ec20baf2131cebc21 Mon Sep 17 00:00:00 2001
From: Stas <sts@abc3.dev>
Date: Tue, 28 May 2024 14:34:26 +0200
Subject: [PATCH 3/3] add logs

---
 lib/supavisor/client_handler.ex | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/lib/supavisor/client_handler.ex b/lib/supavisor/client_handler.ex
index 31546d68..a000f20c 100644
--- a/lib/supavisor/client_handler.ex
+++ b/lib/supavisor/client_handler.ex
@@ -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
         [
@@ -864,6 +866,7 @@ defmodule Supavisor.ClientHandler do
           {:error, reason}
       end
 
+    Logger.info("ClientHandler: Get secrets finished")
     resp
   end