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

@@ -313,6 +313,13 @@ defmodule Supavisor.ClientHandler do
             {:next_state, :auth_md5_wait, %{data | auth_context: auth_context},
              {:timeout, 15_000, :auth_timeout}}
 
+          :auth_query_jit ->
+            auth_context = Auth.create_auth_context(method, secrets, info)
+            :ok = HandlerHelpers.sock_send(sock, Server.password_request())
+
+            {:next_state, :auth_password_wait, %{data | auth_context: auth_context},
+             {:timeout, 15_000, :auth_timeout}}
+
           _scram_method ->
             :ok = HandlerHelpers.sock_send(sock, Server.scram_request())
             auth_context = Auth.create_auth_context(method, secrets, info)
@@ -980,7 +987,9 @@ defmodule Supavisor.ClientHandler do
       method: proxy_type,
       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
     }
 
     %{

lib/supavisor/client_handler/auth.ex

@@ -14,7 +14,7 @@ defmodule Supavisor.ClientHandler.Auth do
   alias Supavisor.{Helpers, Protocol.Server}
   alias Supavisor.ClientHandler.Auth.{MD5Secrets, PasswordSecrets, SASLSecrets}
 
-  @type auth_method :: :password | :auth_query | :auth_query_md5
+  @type auth_method :: :password | :auth_query | :auth_query_md5 | :auth_query_jit
   @type auth_secrets :: {auth_method(), function()}
 
   ## Secret Management
@@ -83,6 +83,57 @@ defmodule Supavisor.ClientHandler.Auth do
       else: {:error, :wrong_password}
   end
 
+  def validate_credentials(: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
+
   ## Challenge Preparation
 
   @doc """
@@ -233,7 +284,8 @@ defmodule Supavisor.ClientHandler.Auth do
     }
   end
 
-  def create_auth_context(method, secrets, info) when method in [:password, :auth_query] do
+  def create_auth_context(method, secrets, info)
+      when method in [:password, :auth_query, :auth_query_jit] do
     %{
       method: method,
       secrets: secrets,
@@ -312,9 +364,10 @@ defmodule Supavisor.ClientHandler.Auth do
 
           with {:ok, secret} <- Helpers.get_user_secret(conn, tenant.auth_query, db_user) do
             auth_type =
-              case secret do
-                %MD5Secrets{} -> :auth_query_md5
-                %SASLSecrets{} -> :auth_query
+              case {tenant.use_jit, secret} do
+                {true, %PasswordSecrets{}} -> :auth_query_jit
+                {_, %MD5Secrets{}} -> :auth_query_md5
+                {_, %SASLSecrets{}} -> :auth_query
               end
 
             {:ok, {auth_type, fn -> secret end}}

lib/supavisor/client_handler/auth/password_secrets.ex

@@ -6,6 +6,8 @@ defmodule Supavisor.ClientHandler.Auth.PasswordSecrets do
 
   @type t :: %__MODULE__{
           user: String.t(),
-          password: String.t()
+          password: String.t(),
+          use_jit: Bool.t(),
+          jit_api_url: String.t()
         }
 end

lib/supavisor/client_handler/error.ex

@@ -157,6 +157,7 @@ defmodule Supavisor.ClientHandler.Error do
     log_message =
       case context do
         :auth_md5_wait -> "Timeout while waiting for MD5 password"
+        :auth_password_wait -> "Timeout while waiting for password"
         :auth_scram_first_wait -> "Timeout while waiting for first SCRAM message"
         :auth_scram_final_wait -> "Timeout while waiting for final SCRAM message"
         _ -> "Authentication timeout"
@@ -274,6 +275,7 @@ defmodule Supavisor.ClientHandler.Error do
   end
 
   defp auth_context_description(:auth_md5_wait), do: "MD5"
+  defp auth_context_description(:auth_password_wait), do: "PASSWORD"
   defp auth_context_description(:auth_scram_first_wait), do: "SCRAM first"
   defp auth_context_description(:auth_scram_final_wait), do: "SCRAM final"
   defp auth_context_description(_), do: "Unknown"

lib/supavisor/db_handler.ex

@@ -251,6 +251,12 @@ defmodule Supavisor.DbHandler do
         Logger.error("DbHandler: Auth error #{inspect(error)}")
         {:stop, :invalid_password, data}
 
+      {:error_response, %{"S" => "FATAL", "C" => "28000"} = error} ->
+        reason = error["M"] || "Authentication failed"
+        handle_authentication_error(data, reason)
+        Logger.error("DbHandler: Auth error #{inspect(error)}")
+        {:stop, :invalid_password, data}
+
       {:error_response, %{"S" => "FATAL", "C" => "3D000"} = error} ->
         Logger.error("DbHandler: Database does not exist: #{inspect(error)}")
         encode_and_forward_error(error, data)
@@ -656,10 +662,17 @@ 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)}")
 
-    {_method, secrets_fn} = data.auth.secrets
+    {method, secrets_fn} = data.auth.secrets
     secrets = secrets_fn.()
 
-    payload = <<secrets.password::binary, 0>>
+    password =
+      if method == :password do
+        secrets.password
+      else
+        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

@@ -435,4 +435,74 @@ defmodule Supavisor.Helpers do
         raise "Invalid boolean value for #{env_var}: #{inspect(value)}. Expected: true, false, 1, or 0"
     end
   end
+
+  @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

@@ -112,6 +112,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 =
@@ -388,11 +391,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

@@ -65,7 +65,9 @@ defmodule Supavisor.SecretChecker do
           ip_version: Supavisor.Helpers.ip_version(tenant.ip_version, tenant.db_host),
           upstream_ssl: tenant.upstream_ssl,
           upstream_verify: tenant.upstream_verify,
-          upstream_tls_ca: Supavisor.Helpers.upstream_cert(tenant.upstream_tls_ca)
+          upstream_tls_ca: Supavisor.Helpers.upstream_cert(tenant.upstream_tls_ca),
+          use_jit: tenant.use_jit,
+          jit_api_url: tenant.jit_api_url
         }
       else
         nil
@@ -117,12 +119,19 @@ defmodule Supavisor.SecretChecker do
     do: Process.send_after(self(), :check, interval + jitter())
 
   def check_secrets(user, %{auth: auth, conn: conn} = state) do
+    alias Supavisor.ClientHandler.Auth
+
+    # 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 =
-          case secret do
-            %Auth.MD5Secrets{} -> :auth_query_md5
-            %Auth.SASLSecrets{} -> :auth_query
+          case {t.use_jit, secret} do
+            {true, %Auth.PasswordSecrets{}} -> :auth_query_jit
+            {_, %Auth.MD5Secrets{}} -> :auth_query_md5
+            {_, %Auth.SASLSecrets{}} -> :auth_query
           end
 
         update_cache =

lib/supavisor/tenants/tenant.ex

@@ -33,6 +33,8 @@ defmodule Supavisor.Tenants.Tenant do
     field(:allow_list, {:array, :string}, default: ["0.0.0.0/0", "::/0"])
     field(:availability_zone, :string)
     field(:feature_flags, :map, default: %{})
+    field(:use_jit, :boolean, default: false)
+    field(:jit_api_url, :string)
 
     has_many(:users, User,
       foreign_key: :tenant_external_id,
@@ -67,7 +69,9 @@ defmodule Supavisor.Tenants.Tenant do
       :client_heartbeat_interval,
       :allow_list,
       :availability_zone,
-      :feature_flags
+      :feature_flags,
+      :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