chore: fixup with rebase

Author: staaldraad

lib/supavisor/client_handler.ex

@@ -620,6 +620,33 @@ defmodule Supavisor.ClientHandler do
     end
   end
 
+  def handle_event(:info, {proto, socket, bin}, :auth_password_wait, data)
+      when proto in @proto do
+    auth_context = data.auth_context
+    {:ok, {ip, _port}} = :inet.peername(socket)
+
+    with {:ok, cls_password} <- Auth.parse_auth_message(bin, auth_context.method),
+         {:ok, key, method} <-
+           Auth.validate_credentials(
+             auth_context.method,
+             auth_context.info.tenant,
+             auth_context.secrets,
+             cls_password,
+             ip
+           ) do
+      handle_auth_success(
+        data.sock,
+        {method, auth_context.secrets},
+        key,
+        cls_password,
+        data
+      )
+    else
+      {:error, reason, _} ->
+        handle_auth_failure(data.sock, reason, data, :auth_password_wait)
+    end
+  end
+
   # SCRAM authentication - waiting for first message
   def handle_event(:info, {proto, _socket, bin}, :auth_scram_first_wait, data)
       when proto in @proto do
@@ -671,7 +698,12 @@ defmodule Supavisor.ClientHandler do
 
   # Authentication timeout handler
   def handle_event(:timeout, :auth_timeout, auth_state, data)
-      when auth_state in [:auth_md5_wait, :auth_scram_first_wait, :auth_scram_final_wait] do
+      when auth_state in [
+             :auth_md5_wait,
+             :auth_password_wait,
+             :auth_scram_first_wait,
+             :auth_scram_final_wait
+           ] do
     handle_auth_failure(data.sock, {:timeout, auth_state}, data, auth_state)
   end
 
@@ -795,9 +827,12 @@ defmodule Supavisor.ClientHandler do
   end
 
   ## Internal functions
-
   defp handle_auth_success(sock, {method, secrets}, client_key, data) do
-    final_secrets = Auth.prepare_final_secrets(secrets, client_key)
+    handle_auth_success(sock, {method, secrets}, client_key, nil, data)
+  end
+
+  defp handle_auth_success(sock, {method, secrets}, client_key, password, data) do
+    final_secrets = Auth.prepare_final_secrets(secrets, client_key, password)
 
     # Only store in TenantCache for pool modes (transaction/session)
     # For proxy mode, secrets are passed directly to DbHandler via data.auth

lib/supavisor/client_handler/auth.ex

@@ -83,39 +83,38 @@ defmodule Supavisor.ClientHandler.Auth do
       else: {:error, :wrong_password}
   end
 
-  def validate_credentials(:auth_query_jit, secrets, password, ip) do
+  def validate_credentials(:auth_query_jit, tenant, 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
+      case Helpers.check_user_has_jit_role(tenant.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, :crypto.hash(:sha256, password), :auth_query_jit}
 
         {:ok, false} ->
           Logger.debug("User token is valid but can't assume this role")
-          {:error, :wrong_password}
+          {:error, :wrong_password, :auth_query_jit}
 
         {:error, :unauthorized_or_forbidden} ->
-          {:error, :wrong_password}
+          {:error, :wrong_password, :auth_query_jit}
 
         {:error, _} ->
           Logger.debug("Unexpected error while calling API")
-          {:error, :wrong_password}
+          {:error, :wrong_password, :auth_query_jit}
       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}
+            {:ok, nil, :auth_query_md5}
           else
-            {:error, :wrong_password}
+            {:error, :wrong_password, :auth_query_md5}
           end
 
         _ ->
@@ -128,8 +127,8 @@ defmodule Supavisor.ClientHandler.Auth do
           stored_key = :crypto.hash(:sha256, client_key)
 
           if :crypto.hash_equals(stored_key, secret.stored_key),
-            do: {:ok, client_key},
-            else: {:error, :wrong_password}
+            do: {:ok, client_key, :auth_query},
+            else: {:error, :wrong_password, :auth_query}
       end
     end
   end
@@ -248,6 +247,9 @@ defmodule Supavisor.ClientHandler.Auth do
 
   def parse_auth_message(bin, _scram_method) do
     case Server.decode_pkt(bin) do
+      {:ok, %{tag: :password_message, payload: {:cleartext_password, cls_password}}, _} ->
+        {:ok, cls_password}
+
       {:ok,
        %{
          tag: :password_message,
@@ -284,8 +286,19 @@ defmodule Supavisor.ClientHandler.Auth do
     }
   end
 
+  @spec create_auth_context(auth_method(), function(), map()) :: map()
+  def create_auth_context(:auth_query_jit, secrets, info) do
+    %{
+      method: :auth_query_jit,
+      secrets: secrets,
+      info: info,
+      cls_password: nil,
+      signatures: nil
+    }
+  end
+
   def create_auth_context(method, secrets, info)
-      when method in [:password, :auth_query, :auth_query_jit] do
+      when method in [:password, :auth_query] do
     %{
       method: method,
       secrets: secrets,
@@ -302,6 +315,14 @@ defmodule Supavisor.ClientHandler.Auth do
     %{auth_context | signatures: signatures}
   end
 
+  @doc """
+  Updates authentication context with new jit information after first exchange.
+  """
+  @spec update_auth_context_with_jit(map(), map()) :: map()
+  def update_auth_context_with_jit(auth_context, cls_password) do
+    %{auth_context | cls_password: cls_password}
+  end
+
   ## Success Response Preparation
 
   @doc """
@@ -327,6 +348,15 @@ defmodule Supavisor.ClientHandler.Auth do
     fn -> Map.put(secrets_fn.(), :client_key, client_key) end
   end
 
+  def prepare_final_secrets(secrets_fn, client_key, password) do
+    fn ->
+      Map.merge(secrets_fn.(), %{
+        client_key: client_key,
+        cls_password: password
+      })
+    end
+  end
+
   ## Private Helpers
 
   @spec fetch_secrets_from_database(Supavisor.id(), map(), String.t()) ::
@@ -365,7 +395,7 @@ defmodule Supavisor.ClientHandler.Auth do
           with {:ok, secret} <- Helpers.get_user_secret(conn, tenant.auth_query, db_user) do
             auth_type =
               case {tenant.use_jit, secret} do
-                {true, %PasswordSecrets{}} -> :auth_query_jit
+                {true, _} -> :auth_query_jit
                 {_, %MD5Secrets{}} -> :auth_query_md5
                 {_, %SASLSecrets{}} -> :auth_query
               end

lib/supavisor/client_handler/auth/password_secrets.ex

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

lib/supavisor/client_handler/auth/sasl_secrets.ex

@@ -1,8 +1,17 @@
 defmodule Supavisor.ClientHandler.Auth.SASLSecrets do
   @moduledoc "Secrets for SCRAM-SHA-256 authentication"
 
-  @derive {Inspect, except: [:client_key, :server_key, :salt, :stored_key]}
-  defstruct [:user, :client_key, :server_key, :digest, :iterations, :salt, :stored_key]
+  @derive {Inspect, except: [:client_key, :server_key, :salt, :stored_key, :cls_password]}
+  defstruct [
+    :user,
+    :client_key,
+    :server_key,
+    :digest,
+    :iterations,
+    :salt,
+    :stored_key,
+    :cls_password
+  ]
 
   @type t :: %__MODULE__{
           user: String.t(),
@@ -11,6 +20,7 @@ defmodule Supavisor.ClientHandler.Auth.SASLSecrets do
           digest: atom(),
           iterations: pos_integer(),
           salt: binary(),
-          stored_key: binary()
+          stored_key: binary(),
+          cls_password: String.t() | nil
         }
 end

lib/supavisor/secret_checker.ex

@@ -129,7 +129,7 @@ defmodule Supavisor.SecretChecker do
       {:ok, secret} ->
         method =
           case {t.use_jit, secret} do
-            {true, %Auth.PasswordSecrets{}} -> :auth_query_jit
+            {true, %Auth.SASLSecrets{}} -> :auth_query_jit
             {_, %Auth.MD5Secrets{}} -> :auth_query_md5
             {_, %Auth.SASLSecrets{}} -> :auth_query
           end