feat(rbac): saml jit provisioning (#119)

## 📝 Description

I have updated our saml login procedure. Now, if the user does not exist
already, and the organization has JIT provisioning enabled within it's
SAML integration, the user will be created on the spot.

renderedtext/tasks#7566

## ✅ Checklist
- [x] I have tested this change
- [x] This change requires documentation update

Author: VeljkoMaksimovic

ee/rbac/lib/internal_api/rbac.pb.ex

@@ -29,6 +29,7 @@ defmodule InternalApi.RBAC.RoleBindingSource do
   field(:ROLE_BINDING_SOURCE_GITLAB, 4)
   field(:ROLE_BINDING_SOURCE_SCIM, 5)
   field(:ROLE_BINDING_SOURCE_INHERITED_FROM_ORG_ROLE, 6)
+  field(:ROLE_BINDING_SOURCE_SAML_JIT, 7)
 end
 
 defmodule InternalApi.RBAC.ListUserPermissionsRequest do

ee/rbac/lib/rbac/front_repo/user.ex

@@ -37,7 +37,7 @@ defmodule Rbac.FrontRepo.User do
     field(:remember_created_at, :utc_datetime)
     field(:visited_at, :utc_datetime)
 
-    field(:creation_source, Ecto.Enum, values: [:okta])
+    field(:creation_source, Ecto.Enum, values: [:okta, :saml_jit])
     field(:single_org_user, :boolean)
     field(:org_id, :binary_id)
     field(:idempotency_token, :string)

ee/rbac/lib/rbac/grpc_servers/rbac_server.ex

@@ -69,7 +69,7 @@ defmodule Rbac.GrpcServers.RbacServer do
           %RBAC.AssignRoleResponse{}
 
         {:error, error_msg} ->
-          IO.puts("Error assigning role: #{error_msg}")
+          Logger.error("Error assigning role: #{error_msg}")
           grpc_error!(:failed_precondition, error_msg)
       end
     end)
@@ -408,6 +408,7 @@ defmodule Rbac.GrpcServers.RbacServer do
           "github" -> :ROLE_BINDING_SOURCE_GITHUB
           "bitbucket" -> :ROLE_BINDING_SOURCE_BITBUCKET
           "okta" -> :ROLE_BINDING_SOURCE_SCIM
+          "saml_jit" -> :ROLE_BINDING_SOURCE_SAML_JIT
           "inherited_from_org_role" -> :ROLE_BINDING_SOURCE_INHERITED_FROM_ORG_ROLE
           _ -> :ROLE_BINDING_SOURCE_UNSPECIFIED
         end

ee/rbac/lib/rbac/okta/saml/api.ex

@@ -99,7 +99,9 @@ defmodule Rbac.Okta.Saml.Api do
   # Okta Authentication
   #
 
-  post "/okta/auth" do
+  post("/okta/auth", do: handle_auth(conn))
+
+  defp handle_auth(conn) do
     alias Rbac.Okta.Integration
     alias Rbac.Okta.Saml.PayloadParser
 
@@ -116,26 +118,42 @@ defmodule Rbac.Okta.Saml.Api do
     consume_uri = "https://#{org_username}.#{domain()}/okta/auth"
     metadata_uri = "https://#{org_username}.#{domain()}"
 
-    with {:ok, integration} <- Integration.find_by_org_id(org_id),
-         {:ok, email} <- PayloadParser.parse(integration, params, consume_uri, metadata_uri),
-         {:ok, okta_user} <- find_okta_user(integration, email),
-         {:ok, user} <- find_user(okta_user),
+    {:ok, integration} = Integration.find_by_org_id(org_id)
+    {:ok, email, attributes} = PayloadParser.parse(integration, params, consume_uri, metadata_uri)
+
+    with {:ok, scim_saml_user} <- find_scim_or_saml_user(integration, email),
+         {:ok, user} <- find_user(scim_saml_user),
          {:ok, user} <- FrontRepo.User.set_remember_timestamp(user) do
-      Watchman.increment("okta_login.success")
+      Watchman.increment("saml_login.success")
 
       Logger.info(
-        "[Okta] User create a session user_id: #{user.id} integration_id: #{integration.id}"
+        "[SAML] User create a session user_id: #{user.id} integration_id: #{integration.id}"
       )
 
       conn
       |> inject_session_cookie(user)
       |> redirect(user)
     else
+      {:error, :scim_saml_user, :not_found} = e ->
+        if integration.jit_provisioning_enabled do
+          {:ok, saml_jit_user} = Rbac.Repo.SamlJitUser.create(integration, email, attributes)
+          :ok = Rbac.Okta.Saml.JitProvisioner.AddUser.run(saml_jit_user)
+
+          conn
+          |> put_resp_content_type("text/plain")
+          |> send_resp(200, "User provisioning, try again in a minute")
+        else
+          raise e
+        end
+
       e ->
-        Watchman.increment("okta_login.failure")
-        Logger.error("Okta auth failed #{inspect(e)}")
-        render_not_found(conn)
+        raise e
     end
+  rescue
+    e ->
+      Watchman.increment("saml_login.failure")
+      Logger.error("SAML auth failed #{inspect(e)}")
+      render_not_found(conn)
   end
 
   defp inject_session_cookie(conn, user) do
@@ -166,16 +184,18 @@ defmodule Rbac.Okta.Saml.Api do
     conn |> put_resp_content_type("text/plain") |> send_resp(404, "Not found")
   end
 
-  defp find_okta_user(integration, email) do
-    case Rbac.Repo.OktaUser.find_by_email(integration, email) do
+  defp find_scim_or_saml_user(integration, email) do
+    with {:error, :not_found} <- Rbac.Repo.OktaUser.find_by_email(integration, email),
+         {:error, :not_found} <- Rbac.Repo.SamlJitUser.find_by_email(integration, email) do
+      {:error, :scim_saml_user, :not_found}
+    else
       {:ok, user} -> {:ok, user}
-      {:error, :not_found} -> {:error, :okta_user, :not_found}
     end
   end
 
-  defp find_user(okta_user) do
-    if okta_user.user_id do
-      case FrontRepo.User.active_user_by_id(okta_user.user_id) do
+  defp find_user(saml_scim_user) do
+    if saml_scim_user.user_id do
+      case FrontRepo.User.active_user_by_id(saml_scim_user.user_id) do
         {:ok, user} -> {:ok, user}
         {:error, :not_found} -> {:error, :user, :not_found}
       end

ee/rbac/lib/rbac/okta/saml/jit_provisioner/add_user.ex

@@ -0,0 +1,104 @@
+# credo:disable-for-this-file
+defmodule Rbac.Okta.Saml.JitProvisioner.AddUser do
+  @moduledoc """
+  Procedure for creating a Semaphore user based on the received
+  okta user payload.
+
+  The assumed state of the system:
+  - We have an okta_user record in the database
+  - The okta_user.state is :pending
+  - The okta_user.user_id is nil
+  """
+
+  require Logger
+
+  alias Rbac.Repo.{SamlJitUser, RbacRole}
+  alias Rbac.RoleBindingIdentification
+  alias Rbac.RoleManagement
+
+  @role_name "Member"
+
+  def run(saml_jit_user) do
+    idempotency_token = "okta-user-#{saml_jit_user.id}"
+    email = saml_jit_user.email
+    name = SamlJitUser.construct_name(saml_jit_user)
+
+    user_params = %{
+      email: email,
+      name: name,
+      idempotency_token: idempotency_token,
+      creation_source: :saml_jit,
+      single_org_user: true,
+      org_id: saml_jit_user.org_id
+    }
+
+    Logger.info("Provisioning #{saml_jit_user.id}")
+
+    with(
+      {:ok, user} <- find_or_create_user(idempotency_token, user_params),
+      {:ok, saml_jit_user} <- SamlJitUser.connect_user(saml_jit_user, user.id),
+      :ok <- assign_role(user.id, saml_jit_user.org_id, @role_name),
+      {:ok, _saml_jit_user} <- SamlJitUser.mark_as_processed(saml_jit_user)
+    ) do
+      Logger.info("Provisioning #{saml_jit_user.id} done.")
+
+      :ok
+    else
+      err ->
+        log_provisioning_error(saml_jit_user, err)
+        err
+    end
+  end
+
+  defp find_or_create_user(idempotency_token, user_params) do
+    case find_user_by_idempotency_token(idempotency_token) do
+      nil ->
+        case Rbac.Store.RbacUser.fetch_by_email(user_params.email) do
+          {:error, :not_found} ->
+            Logger.info("[Sam lJIT Provisioner] Creating new user #{inspect(user_params)}")
+            Rbac.User.Actions.create(user_params)
+
+          {:ok, user} ->
+            Logger.info(
+              "[Saml JIT Provisioner] Adding idempotency token to existing user #{inspect(user_params)}"
+            )
+
+            Rbac.Store.User.Front.add_idempotency_token(user.id, idempotency_token)
+            {:ok, user}
+        end
+
+      user ->
+        {:ok, user}
+    end
+  end
+
+  defp find_user_by_idempotency_token(idempotency_token) do
+    case Rbac.Store.User.Front.find_by_idempotency_token(idempotency_token) do
+      {:error, :not_found} ->
+        nil
+
+      {:ok, user} ->
+        Rbac.Store.RbacUser.fetch(user.id)
+    end
+  end
+
+  defp assign_role(user_id, org_id, role_name) do
+    if RoleManagement.user_part_of_org?(user_id, org_id) do
+      :ok
+    else
+      with {:ok, rbi} <- RoleBindingIdentification.new(user_id: user_id, org_id: org_id),
+           {:ok, role} <- RbacRole.get_role_by_name(role_name, "org_scope", org_id) do
+        {:ok, nil} = RoleManagement.assign_role(rbi, role.id, :saml_jit)
+
+        Rbac.Events.UserJoinedOrganization.publish(user_id, org_id)
+        :ok
+      end
+    end
+  end
+
+  defp log_provisioning_error(okta_user, err) do
+    inspects = "okta user: #{inspect(okta_user)} error: #{inspect(err)}"
+
+    Logger.error("SamlJIT Provisioner: Failed to provision #{inspects}")
+  end
+end

ee/rbac/lib/rbac/okta/saml/payload_parser.ex

@@ -14,10 +14,10 @@ defmodule Rbac.Okta.Saml.PayloadParser do
          {:ok, payload} <- extract_saml_response(params),
          {:ok, decoded} <- decode_payload(payload),
          {:ok, assertion} <- validate_assertion(decoded, sp) do
-      subject = esaml_assertion(assertion, :subject)
-      email = esaml_subject(subject, :name)
+      email = esaml_assertion(assertion, :subject) |> esaml_subject(:name)
+      attributes = esaml_assertion(assertion, :attributes) |> construct_attributes_map()
 
-      {:ok, to_string(email)}
+      {:ok, to_string(email), attributes}
     end
   end
 
@@ -42,6 +42,13 @@ defmodule Rbac.Okta.Saml.PayloadParser do
       {:error, :invalid_xml, e}
   end
 
+  defp construct_attributes_map(attributes) do
+    Enum.reduce(attributes, %{}, fn {name, value}, acc ->
+      value = to_string(value)
+      Map.update(acc, name, [value], &(&1 ++ [value]))
+    end)
+  end
+
   defp validate_assertion(saml, sp) do
     :esaml_sp.validate_assertion(saml, sp)
   end

ee/rbac/lib/rbac/repo/saml_jit_user.ex

@@ -0,0 +1,89 @@
+defmodule Rbac.Repo.SamlJitUser do
+  use Rbac.Repo.Schema
+  alias Rbac.Repo
+  import Ecto.Query, only: [where: 3]
+
+  @timestamps_opts [type: :utc_datetime]
+
+  @required_fields [
+    :org_id,
+    :integration_id,
+    :attributes,
+    :state,
+    :email
+  ]
+
+  @updatable_fields [
+    :attributes,
+    :state,
+    :email,
+    :updated_at,
+    :user_id
+  ]
+
+  schema "saml_jit_users" do
+    belongs_to(:integration, Repo.OktaIntegration)
+
+    field(:org_id, :binary_id)
+    field(:attributes, :map)
+    field(:state, Ecto.Enum, values: [:pending, :processed])
+    field(:user_id, :binary_id)
+    field(:email, :string)
+
+    timestamps()
+  end
+
+  def create(integration, email, attributes) do
+    new(integration, email, attributes)
+    |> changeset()
+    |> Rbac.Repo.insert()
+  end
+
+  def find_by_email(integration, email) do
+    __MODULE__
+    |> where([u], u.integration_id == ^integration.id and u.email == ^email)
+    |> Repo.one()
+    |> case do
+      nil -> {:error, :not_found}
+      user -> {:ok, user}
+    end
+  end
+
+  def connect_user(%__MODULE__{} = user, user_id) do
+    changeset(user, %{user_id: user_id}) |> Repo.update()
+  end
+
+  def construct_name(%__MODULE__{} = user) do
+    name = extract_attribute(user, "firstName") <> " " <> extract_attribute(user, "lastName")
+
+    if String.trim(name) == "" do
+      user.email |> String.split("@") |> List.first()
+    else
+      name
+    end
+  end
+
+  def mark_as_processed(%__MODULE__{} = user) do
+    changeset(user, %{state: :processed}) |> Rbac.Repo.update()
+  end
+
+  defp new(integration, email, attributes) do
+    %__MODULE__{
+      integration_id: integration.id,
+      org_id: integration.org_id,
+      attributes: attributes,
+      email: email,
+      state: :pending
+    }
+  end
+
+  defp changeset(%__MODULE__{} = user, params \\ %{}) do
+    user
+    |> cast(params, @updatable_fields)
+    |> validate_required(@required_fields)
+  end
+
+  defp extract_attribute(%__MODULE__{} = user, name) do
+    Map.get(user.attributes, name, [""]) |> List.first()
+  end
+end

ee/rbac/lib/rbac/repo/subject_role_binding.ex

@@ -3,7 +3,7 @@ defmodule Rbac.Repo.SubjectRoleBinding do
   alias Rbac.Repo.{RbacRole, Subject}
   import Ecto.Query, only: [where: 3]
 
-  @binding_sources ~w(github bitbucket gitlab manually_assigned okta inherited_from_org_role)a
+  @binding_sources ~w(github bitbucket gitlab manually_assigned okta inherited_from_org_role saml_jit)a
 
   schema "subject_role_bindings" do
     belongs_to(:role, RbacRole)

ee/rbac/priv/repo/migrations/20250903231134_create_saml_jit_users.exs

@@ -0,0 +1,18 @@
+defmodule Rbac.Repo.Migrations.CreateSamlJitUsers do
+  use Ecto.Migration
+
+  def change do
+    create table("saml_jit_users") do
+      add(:integration_id, :binary_id, null: false)
+      add(:org_id, :binary_id, null: false)
+      add(:attributes, :jsonb, null: false)
+      add(:state, :string, null: false)
+      add(:user_id, :uuid)
+      add(:email, :string)
+
+      timestamps()
+    end
+
+    create(unique_index(:saml_jit_users, [:integration_id, :email]))
+  end
+end

ee/rbac/priv/repo/migrations/20250903233405_add_saml_jit_to_role_binding_type_enum.exs

@@ -0,0 +1,22 @@
+defmodule Rbac.Repo.Migrations.AddSamlJitToRoleBindingEnum do
+  use Ecto.Migration
+
+  ###
+  ### This query efectively adds one more value to enum type
+  ###
+  ### There is a dedicated ALTER TYPE _ ADD VALUE command, but it can not be executed
+  ### Within the transaction, and ecto executes commands inside the transaction, so it can't
+  ### be used here, and we have this long transaction to achive the same thing
+  ###
+  @sql_commands_for_altering_enum_type [
+    "ALTER TYPE role_binding_scope RENAME TO deprecated;",
+    "CREATE TYPE role_binding_scope AS ENUM ('github', 'bitbucket', 'gitlab', 'manually_assigned', 'okta', 'inherited_from_org_role', 'saml_jit');",
+    "ALTER TABLE subject_role_bindings
+      ALTER COLUMN binding_source TYPE role_binding_scope USING binding_source::text::role_binding_scope;",
+    "DROP TYPE deprecated;"
+  ]
+
+  def change do
+    Enum.each(@sql_commands_for_altering_enum_type, &execute/1)
+  end
+end