fix: check cache for secret

Only using data.auth.password would reuse the password for the
is_manager user when trying to authenticate against the upstream.

The first request, from the is_manager to the upstream should use the
password, while the subsequent connection with the authenticating user
should use secret.

feat: support e2e password flow

chore: check password against auth_query

chore: set client_key

Sets the client_key when using password auth (jit) and the incomming
password is a valid password that matches the scram-sha-256. This allows
supporting the case where the tenant has switched back to scram-sha-256
but pooler is still checking for cleartext password.

chore: todo - api implementation

feat: support api lookup

chore: ensure user secret caching works with jit

chore: add jit expire checks

chore: cleanup credo warnings

chore: store client_key for cleartext auth exchange

If the cleartext auth exchange was for a normal password (not PAT/JWT),
we will have a valid client_key that was calculated for the
scram-sha-256. Ensure this is always present in secrets, along with the
clear password if we have it. This allows for the case where the
upstream server either stops using JIT and reverts to scram-sha-256. Or
has a pg_hba entry for the specific user to use scram-sha-256 rather
than pam. This allows both JIT users and non-JIT users to exist in the
same tenant.

chore: support md5 on JIT, lazy eval secrets

fix: migration

fix: terminate does not need additional message

Update lib/supavisor/helpers.ex

Co-authored-by: felipe stival <14948182+v0idpwn@users.noreply.github.com>

chore: start adding more tests

test: add jit api tester

chore: use simplified API with role, rhost

Author: staaldraad

lib/supavisor/client_handler.ex

@@ -357,8 +357,8 @@ defmodule Supavisor.ClientHandler do
           :no_retry -> handle_auth_failure(sock, reason, data)
         end
 
-      {:ok, client_key} ->
-        handle_auth_success(sock, {method, secrets}, client_key, data)
+      {:ok, client_key, password} ->
+        handle_auth_success(sock, {method, secrets}, client_key, password, data)
     end
   end
 
@@ -664,6 +664,7 @@ defmodule Supavisor.ClientHandler do
   def terminate(reason, state, _data) do
     Logger.metadata(state: state)
     Logger.debug("ClientHandler: socket closed with reason #{inspect(reason)}")
+
     :ok
   end
 
@@ -681,7 +682,7 @@ defmodule Supavisor.ClientHandler do
       # Retry authentication with fresh secrets
       case handle_exchange(sock, {method2, secrets2}) do
         {:ok, client_key} ->
-          {:retry_success, handle_auth_success(sock, {method2, secrets2}, client_key, data)}
+          {:retry_success, handle_auth_success(sock, {method2, secrets2}, client_key, nil, data)}
 
         {:error, retry_reason} ->
           Logger.error("ClientHandler: Authentication retry failed: #{inspect(retry_reason)}")
@@ -747,11 +748,17 @@ defmodule Supavisor.ClientHandler do
     Cachex.update(Supavisor.Cache, {:secrets, tenant, user}, {:cached, value})
   end
 
-  defp handle_auth_success(sock, {method, secrets}, client_key, data) do
+  defp handle_auth_success(sock, {method, secrets}, client_key, password, data) do
     final_secrets =
-      if client_key,
-        do: fn -> Map.put(secrets.(), :client_key, client_key) end,
-        else: secrets
+      fn ->
+        Enum.reduce(
+          [client_key: client_key, cls_password: password],
+          secrets.(),
+          fn {k, v}, acc ->
+            if v != nil, do: Map.put(acc, k, v), else: acc
+          end
+        )
+      end
 
     Logger.debug("ClientHandler: Exchange success")
     :ok = HandlerHelpers.sock_send(sock, Server.authentication_ok())
@@ -793,7 +800,7 @@ defmodule Supavisor.ClientHandler do
   end
 
   @spec handle_exchange(Supavisor.sock(), {atom(), fun()}) ::
-          {:ok, binary() | nil}
+          {:ok, binary() | nil, binary() | nil}
           | {:error, :wrong_password | {:timeout, String.t()} | {:unexpected_message, term()}}
   def handle_exchange({_, socket} = sock, {:auth_query_md5 = method, secrets}) do
     salt = :crypto.strong_rand_bytes(4)
@@ -805,7 +812,24 @@ defmodule Supavisor.ClientHandler do
             payload: {:md5, client_md5}
           }, _} <- receive_next(socket, "Timeout while waiting for the md5 exchange"),
          {:ok, key} <- authenticate_exchange(method, client_md5, secrets.().secret, salt) do
-      {:ok, key}
+      {:ok, key, nil}
+    else
+      {:error, message} -> {:error, message}
+      other -> {:error, "Unexpected message #{inspect(other)}"}
+    end
+  end
+
+  def handle_exchange({_, socket} = sock, {:auth_query_jit = method, secrets}) do
+    :ok = HandlerHelpers.sock_send(sock, Server.password_request())
+    {:ok, {ip, _port}} = :inet.peername(socket)
+
+    with {:ok,
+          %{
+            tag: :password_message,
+            payload: {:cleartext_password, password}
+          }, _} <- receive_next(socket, "Timeout while waiting for the password exchange"),
+         {:ok, key} <- authenticate_exchange(method, secrets, password, ip) do
+      {:ok, key, password}
     else
       {:error, message} -> {:error, message}
       other -> {:error, "Unexpected message #{inspect(other)}"}
@@ -839,7 +863,7 @@ defmodule Supavisor.ClientHandler do
          {:ok, key} <- authenticate_exchange(method, secrets, signatures, p) do
       message = "v=#{Base.encode64(signatures.server)}"
       :ok = HandlerHelpers.sock_send(sock, Server.exchange_message(:final, message))
-      {:ok, key}
+      {:ok, key, nil}
     else
       {:error, message} -> {:error, message}
       other -> {:error, "Unexpected message #{inspect(other)}"}
@@ -886,6 +910,57 @@ defmodule Supavisor.ClientHandler do
       else: {:error, :wrong_password}
   end
 
+  defp authenticate_exchange(:auth_query_jit, secrets, password, ip) do
+    # check if incomming password looks like PAT or a JWT
+    # otherwise handle as password,
+    secret = secrets.()
+
+    if Helpers.token_matches?(password) do
+      Logger.debug("Looks like a PAT/JWT - #{inspect(secret.jit_api_url)}")
+      rhost = ip |> :inet.ntoa() |> to_string()
+
+      case Helpers.check_user_has_jit_role(secret.jit_api_url, password, secret.user, rhost) do
+        {:ok, true} ->
+          # set a fake client_key incase upstream switches away from pam mid auth
+          {:ok, :crypto.hash(:sha256, password)}
+
+        {:ok, false} ->
+          Logger.debug("User token is valid but can't assume this role")
+          {:error, :wrong_password}
+
+        {:error, :unauthorized_or_forbidden} ->
+          {:error, :wrong_password}
+
+        {:error, _} ->
+          Logger.debug("Unexpected error while calling API")
+          {:error, :wrong_password}
+      end
+    else
+      # match against the scram-sha-256 / md5 we have from auth_query
+      case secret.digest do
+        :md5 ->
+          if Helpers.md5([password, secret.user]) == secret.secret do
+            {:ok, nil}
+          else
+            {:error, :wrong_password}
+          end
+
+        _ ->
+          salt = Base.decode64!(secret.salt)
+
+          salted_password =
+            :crypto.pbkdf2_hmac(:sha256, password, salt, secret.iterations, 32)
+
+          client_key = :crypto.mac(:hmac, :sha256, salted_password, "Client Key")
+          stored_key = :crypto.hash(:sha256, client_key)
+
+          if :crypto.hash_equals(stored_key, secret.stored_key),
+            do: {:ok, client_key},
+            else: {:error, :wrong_password}
+      end
+    end
+  end
+
   @spec db_checkout(:write | :read | :both, :on_connect | :on_query, map) ::
           {pid, pid, Supavisor.sock()} | nil
   defp db_checkout(_, _, %{mode: mode, db_pid: {pool, db_pid, db_sock}})
@@ -949,7 +1024,9 @@ defmodule Supavisor.ClientHandler do
       require_user: info.tenant.require_user,
       upstream_ssl: info.tenant.upstream_ssl,
       upstream_tls_ca: info.tenant.upstream_tls_ca,
-      upstream_verify: info.tenant.upstream_verify
+      upstream_verify: info.tenant.upstream_verify,
+      use_jit: info.tenant.use_jit,
+      jit_api_url: info.tenant.jit_api_url
     }
 
     %{
@@ -1041,8 +1118,18 @@ defmodule Supavisor.ClientHandler do
 
           resp =
             with {:ok, secret} <- Helpers.get_user_secret(conn, tenant.auth_query, db_user) do
-              t = if secret.digest == :md5, do: :auth_query_md5, else: :auth_query
-              {:ok, {t, fn -> Map.put(secret, :alias, user.db_user_alias) end}}
+              t =
+                case {tenant.use_jit, secret.digest} do
+                  {true, _} -> :auth_query_jit
+                  {_, :md5} -> :auth_query_md5
+                  _ -> :auth_query
+                end
+
+              {:ok,
+               {t,
+                fn ->
+                  Map.merge(secret, %{alias: user.db_user_alias, jit_api_url: tenant.jit_api_url})
+                end}}
             end
 
           Logger.info("ClientHandler: Get secrets finished")

lib/supavisor/db_handler.ex

@@ -183,6 +183,11 @@ defmodule Supavisor.DbHandler do
         Logger.error("DbHandler: Auth error #{inspect(reason)}")
         {:stop, :invalid_password, data}
 
+      {:error_response, ["SFATAL", "VFATAL", "C28000", reason, _, _, _]} ->
+        handle_authentication_error(data, reason)
+        Logger.error("DbHandler: Auth error #{inspect(reason)}")
+        {:stop, :invalid_password, data}
+
       {:error_response, error} ->
         Logger.error("DbHandler: Error auth response #{inspect(error)}")
         {:stop, {:encode_and_forward, error}}
@@ -663,7 +668,14 @@ defmodule Supavisor.DbHandler do
   defp handle_auth_pkts(%{payload: :authentication_cleartext_password} = dec_pkt, _, data) do
     Logger.debug("DbHandler: dec_pkt, #{inspect(dec_pkt, pretty: true)}")
 
-    payload = <<data.auth.password.()::binary, 0>>
+    password =
+      if data.auth.method == :password do
+        data.auth.password.()
+      else
+        data.auth.secrets.().cls_password
+      end
+
+    payload = <<password::binary, 0>>
     bin = [?p, <<IO.iodata_length(payload) + 4::signed-32>>, payload]
     :ok = HandlerHelpers.sock_send(data.sock, bin)
     :authentication_cleartext

lib/supavisor/helpers.ex

@@ -399,4 +399,74 @@ defmodule Supavisor.Helpers do
   """
   @spec mb_to_words(number()) :: pos_integer()
   def mb_to_words(mb), do: round(mb * 1_048_576 / :erlang.system_info(:wordsize))
+
+  @doc """
+  Checks if a token matches either a PAT or JWT.
+
+  ## Parameters
+
+  - `token`: The token string
+  """
+  def token_matches?(<<"sbp_", _rest::binary>>), do: true
+
+  def token_matches?(token) do
+    case String.split(token, ".") do
+      ["eyJ" <> _, "eyJ" <> _, _sig] ->
+        true
+
+      _ ->
+        false
+    end
+  end
+
+  @doc """
+  Makes an HTTPS GET request to `url` with a Bearer `token`
+  and checks if the given `role` is present in the returned `user_roles` list.
+
+  Returns:
+    - `{:ok, true}` if role is present
+    - `{:ok, false}` if role is absent
+    - `{:error, :unauthorized}` for 401
+    - `{:error, :forbidden}` for 403
+    - `{:error, {:unexpected_status, status}}` for other HTTP codes
+  """
+  def check_user_has_jit_role(url, token, role \\ "postgres", rhost, opts \\ []) do
+    opts =
+      Keyword.merge(opts,
+        headers: [
+          authorization: "Bearer #{token}",
+          "content-type": "application/json"
+        ]
+      )
+
+    body =
+      %{
+        rhost: rhost,
+        role: role
+      }
+      |> Jason.encode!()
+
+    response =
+      Req.post!(
+        url,
+        opts
+        |> Keyword.put(:body, body)
+      )
+
+    case response.status do
+      200 ->
+        %{"user_role" => %{"role" => urole}} = response.body
+
+        # double check server response
+        has_valid_role? = urole == role
+
+        {:ok, has_valid_role?}
+
+      status when status in [401, 403] ->
+        {:error, :unauthorized_or_forbidden}
+
+      status ->
+        {:error, {:unexpected_status, status}}
+    end
+  end
 end

lib/supavisor/protocol/server.ex

@@ -73,6 +73,9 @@ defmodule Supavisor.Protocol.Server do
   @spec md5_request(<<_::32>>) :: iodata()
   def md5_request(salt), do: [<<?R, 12::32, 5::32>>, salt]
 
+  @spec password_request() :: iodata()
+  def password_request, do: [<<?R, 8::32, 3::32>>]
+
   @spec exchange_first_message(binary, binary | boolean, pos_integer) :: binary
   def exchange_first_message(nonce, salt \\ false, iterations \\ 4096) do
     server_nonce =
@@ -329,11 +332,17 @@ defmodule Supavisor.Protocol.Server do
   end
 
   @spec decode_payload(:password_message, binary()) ::
-          {:first_msg_response, map()} | :undefined
+          {:first_msg_response, map()} | {:cleartext_password, binary()} | :undefined
   defp decode_payload(:password_message, bin) do
     case kv_to_map(bin) do
-      {:ok, map} -> {:first_msg_response, map}
-      {:error, _} -> :undefined
+      {:ok, map} ->
+        {:first_msg_response, map}
+
+      {:error, _} ->
+        case :binary.split(bin, <<0>>) do
+          [password, ""] -> {:cleartext_password, password}
+          _ -> :undefined
+        end
     end
   end
 

lib/supavisor/secret_checker.ex

@@ -101,10 +101,31 @@ defmodule Supavisor.SecretChecker do
     do: Process.send_after(self(), :check, interval)
 
   def check_secrets(user, %{auth: auth, conn: conn} = state) do
+    # get tenant information so that we can check configuration information
+    # that affects secrets, like :use_jit
+    t = Supavisor.Tenants.get_tenant_cache(state.tenant, state.auth.sni_hostname)
+
     case Helpers.get_user_secret(conn, auth.auth_query, user) do
       {:ok, secret} ->
-        method = if secret.digest == :md5, do: :auth_query_md5, else: :auth_query
-        secrets = Map.put(secret, :alias, auth.alias)
+        # when jit is in use, make sure to keep the jit_api_url in the cache, and
+        # set the method to also be otherwise
+        # we'll try use scram-sha-256 or md5 even though the upstream server is expecting
+        # jit. Something that is configured in the tenant.
+        # make sure to bring across the jit_api_url too, which will never be in the secrets
+        # returned by the db
+        method =
+          cond do
+            t.use_jit -> :auth_query_jit
+            secret.digest == :md5 -> :auth_query_md5
+            true -> :auth_query
+          end
+
+        secrets =
+          if t.use_jit do
+            Map.merge(secret, %{alias: auth.alias, jit_api_url: t.jit_api_url})
+          else
+            Map.put(secret, :alias, auth.alias)
+          end
 
         update_cache =
           case Cachex.get(Supavisor.Cache, state.key) do

lib/supavisor/tenants/tenant.ex

@@ -32,6 +32,8 @@ defmodule Supavisor.Tenants.Tenant do
     field(:client_heartbeat_interval, :integer, default: 60)
     field(:allow_list, {:array, :string}, default: ["0.0.0.0/0", "::/0"])
     field(:availability_zone, :string)
+    field(:use_jit, :boolean, default: false)
+    field(:jit_api_url, :string)
 
     has_many(:users, User,
       foreign_key: :tenant_external_id,
@@ -65,7 +67,9 @@ defmodule Supavisor.Tenants.Tenant do
       :client_idle_timeout,
       :client_heartbeat_interval,
       :allow_list,
-      :availability_zone
+      :availability_zone,
+      :use_jit,
+      :jit_api_url
     ])
     |> check_constraint(:upstream_ssl, name: :upstream_constraints, prefix: "_supavisor")
     |> check_constraint(:upstream_verify, name: :upstream_constraints, prefix: "_supavisor")

priv/repo/migrations/20250811122429_add_tenant_use_it.exs

@@ -0,0 +1,10 @@
+defmodule Supavisor.Repo.Migrations.AddJitConfig do
+  use Ecto.Migration
+
+  def change do
+    alter table("tenants", prefix: "_supavisor") do
+      add(:use_jit, :boolean, default: false)
+      add(:jit_api_url, :string)
+    end
+  end
+end