feat: sec 488 support jit access in supavisor

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)
@@ -613,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
@@ -664,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
 
@@ -788,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
@@ -980,7 +1022,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,56 @@ defmodule Supavisor.ClientHandler.Auth do
       else: {:error, :wrong_password}
   end
 
+  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
+      rhost = ip |> :inet.ntoa() |> to_string()
+
+      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), :auth_query_jit}
+
+        {:ok, false} ->
+          Logger.debug("User token is valid but can't assume this role")
+          {:error, :wrong_password, :auth_query_jit}
+
+        {:error, :unauthorized_or_forbidden} ->
+          {:error, :wrong_password, :auth_query_jit}
+
+        {:error, _} ->
+          Logger.debug("Unexpected error while calling API")
+          {: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, :auth_query_md5}
+          else
+            {:error, :wrong_password, :auth_query_md5}
+          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, :auth_query},
+            else: {:error, :wrong_password, :auth_query}
+      end
+    end
+  end
+
   ## Challenge Preparation
 
   @doc """
@@ -197,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,
@@ -233,7 +286,19 @@ defmodule Supavisor.ClientHandler.Auth do
     }
   end
 
-  def create_auth_context(method, secrets, info) when method in [:password, :auth_query] do
+  @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] do
     %{
       method: method,
       secrets: secrets,
@@ -250,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 """
@@ -275,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()) ::
@@ -312,9 +394,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, _} -> :auth_query_jit
+                {_, %MD5Secrets{}} -> :auth_query_md5
+                {_, %SASLSecrets{}} -> :auth_query
               end
 
             {:ok, {auth_type, fn -> secret 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/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