feat(front): Role mapping for SAML JIT provisioned users (#123)

## 📝 Description
<!-- Describe your changes in detail -->

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

Author: hamir-suspect

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

@@ -52,37 +52,59 @@ defmodule InternalApi.Okta.GenerateScimTokenResponse do
   field(:token, 1, type: :string)
 end
 
-defmodule InternalApi.Okta.SetUpGroupMappingRequest do
+defmodule InternalApi.Okta.SetUpMappingRequest do
   @moduledoc false
 
   use Protobuf, syntax: :proto3, protoc_gen_elixir_version: "0.13.0"
 
   field(:org_id, 1, type: :string, json_name: "orgId")
   field(:default_role_id, 2, type: :string, json_name: "defaultRoleId")
-  field(:mappings, 3, repeated: true, type: InternalApi.Okta.GroupMapping)
+
+  field(:group_mapping, 3,
+    repeated: true,
+    type: InternalApi.Okta.GroupMapping,
+    json_name: "groupMapping"
+  )
+
+  field(:role_mapping, 4,
+    repeated: true,
+    type: InternalApi.Okta.RoleMapping,
+    json_name: "roleMapping"
+  )
 end
 
-defmodule InternalApi.Okta.SetUpGroupMappingResponse do
+defmodule InternalApi.Okta.SetUpMappingResponse do
   @moduledoc false
 
   use Protobuf, syntax: :proto3, protoc_gen_elixir_version: "0.13.0"
 end
 
-defmodule InternalApi.Okta.DescribeGroupMappingRequest do
+defmodule InternalApi.Okta.DescribeMappingRequest do
   @moduledoc false
 
   use Protobuf, syntax: :proto3, protoc_gen_elixir_version: "0.13.0"
 
   field(:org_id, 1, type: :string, json_name: "orgId")
 end
 
-defmodule InternalApi.Okta.DescribeGroupMappingResponse do
+defmodule InternalApi.Okta.DescribeMappingResponse do
   @moduledoc false
 
   use Protobuf, syntax: :proto3, protoc_gen_elixir_version: "0.13.0"
 
   field(:default_role_id, 1, type: :string, json_name: "defaultRoleId")
-  field(:mappings, 2, repeated: true, type: InternalApi.Okta.GroupMapping)
+
+  field(:group_mapping, 2,
+    repeated: true,
+    type: InternalApi.Okta.GroupMapping,
+    json_name: "groupMapping"
+  )
+
+  field(:role_mapping, 3,
+    repeated: true,
+    type: InternalApi.Okta.RoleMapping,
+    json_name: "roleMapping"
+  )
 end
 
 defmodule InternalApi.Okta.GroupMapping do
@@ -94,6 +116,15 @@ defmodule InternalApi.Okta.GroupMapping do
   field(:okta_group_id, 2, type: :string, json_name: "oktaGroupId")
 end
 
+defmodule InternalApi.Okta.RoleMapping do
+  @moduledoc false
+
+  use Protobuf, syntax: :proto3, protoc_gen_elixir_version: "0.13.0"
+
+  field(:semaphore_role_id, 1, type: :string, json_name: "semaphoreRoleId")
+  field(:okta_role_id, 2, type: :string, json_name: "oktaRoleId")
+end
+
 defmodule InternalApi.Okta.ListRequest do
   @moduledoc false
 
@@ -160,16 +191,12 @@ defmodule InternalApi.Okta.Okta.Service do
 
   rpc(:Destroy, InternalApi.Okta.DestroyRequest, InternalApi.Okta.DestroyResponse)
 
-  rpc(
-    :SetUpGroupMapping,
-    InternalApi.Okta.SetUpGroupMappingRequest,
-    InternalApi.Okta.SetUpGroupMappingResponse
-  )
+  rpc(:SetUpMapping, InternalApi.Okta.SetUpMappingRequest, InternalApi.Okta.SetUpMappingResponse)
 
   rpc(
-    :DescribeGroupMapping,
-    InternalApi.Okta.DescribeGroupMappingRequest,
-    InternalApi.Okta.DescribeGroupMappingResponse
+    :DescribeMapping,
+    InternalApi.Okta.DescribeMappingRequest,
+    InternalApi.Okta.DescribeMappingResponse
   )
 end
 

ee/rbac/lib/rbac/grpc_servers/okta_server.ex

@@ -10,9 +10,10 @@ defmodule Rbac.GrpcServers.OktaServer do
     GenerateScimTokenResponse,
     ListResponse,
     ListUsersResponse,
-    SetUpGroupMappingResponse,
-    DescribeGroupMappingResponse,
-    GroupMapping
+    SetUpMappingResponse,
+    DescribeMappingResponse,
+    GroupMapping,
+    RoleMapping
   }
 
   @manage_okta_permission "organization.okta.manage"
@@ -128,69 +129,89 @@ defmodule Rbac.GrpcServers.OktaServer do
     end)
   end
 
-  def set_up_group_mapping(req, _stream) do
-    observe("set_up_group_mapping", fn ->
+  def set_up_mapping(req, _stream) do
+    observe("set_up_mapping", fn ->
       validate_uuid!(req.org_id)
 
-      mappings =
-        Enum.map(req.mappings, fn mapping ->
+      group_mappings =
+        Enum.map(req.group_mapping, fn mapping ->
           %{
             idp_group_id: mapping.okta_group_id,
             semaphore_group_id: mapping.semaphore_group_id
           }
         end)
 
-      case Rbac.Okta.IdpGroupMapping.create_or_update(req.org_id, mappings, req.default_role_id) do
+      role_mappings =
+        Enum.map(req.role_mapping, fn mapping ->
+          %{
+            idp_role_id: mapping.okta_role_id,
+            semaphore_role_id: mapping.semaphore_role_id
+          }
+        end)
+
+      case Rbac.Okta.IdpGroupMapping.create_or_update(
+             req.org_id,
+             group_mappings,
+             role_mappings,
+             req.default_role_id
+           ) do
         {:ok, _mapping} ->
-          Logger.info("Group mappings created/updated for org #{req.org_id}")
-          %SetUpGroupMappingResponse{}
+          Logger.info("Group and role mappings created/updated for org #{req.org_id}")
+          %SetUpMappingResponse{}
 
         {:error, %Ecto.Changeset{} = changeset} ->
           Logger.error(
-            "Error while setting up group mappings for org #{req.org_id}: #{inspect(changeset)}"
+            "Error while setting up mappings for org #{req.org_id}: #{inspect(changeset)}"
           )
 
           grpc_error!(
             :failed_precondition,
-            "Failed to save group mappings: Invalid"
+            "Failed to save mappings: Invalid"
           )
 
         error ->
-          Logger.error("Unknown error while setting up group mappings: #{inspect(error)}")
-          grpc_error!(:unknown, "Unknown error while setting up group mappings")
+          Logger.error("Unknown error while setting up mappings: #{inspect(error)}")
+          grpc_error!(:unknown, "Unknown error while setting up mappings")
       end
     end)
   end
 
-  def describe_group_mapping(req, _stream) do
-    observe("list_group_mappings", fn ->
+  def describe_mapping(req, _stream) do
+    observe("describe_mapping", fn ->
       validate_uuid!(req.org_id)
 
       case Rbac.Okta.IdpGroupMapping.get_for_organization(req.org_id) do
-        {:ok, idp_group_mapping} ->
+        {:ok, idp_mapping} ->
           # Convert from our internal format to protobuf messages
           group_mapping =
-            Enum.map(idp_group_mapping.group_mapping, fn mapping ->
+            Enum.map(idp_mapping.group_mapping, fn mapping ->
               %GroupMapping{
                 okta_group_id: mapping.idp_group_id,
                 semaphore_group_id: mapping.semaphore_group_id
               }
             end)
 
-          %DescribeGroupMappingResponse{
-            mappings: group_mapping,
-            default_role_id: idp_group_mapping.default_role_id
+          role_mapping =
+            Enum.map(idp_mapping.role_mapping || [], fn mapping ->
+              %RoleMapping{
+                okta_role_id: mapping.idp_role_id,
+                semaphore_role_id: mapping.semaphore_role_id
+              }
+            end)
+
+          %DescribeMappingResponse{
+            group_mapping: group_mapping,
+            role_mapping: role_mapping,
+            default_role_id: idp_mapping.default_role_id
           }
 
         {:error, :not_found} ->
-          %DescribeGroupMappingResponse{mappings: []}
+          %DescribeMappingResponse{group_mapping: [], role_mapping: []}
 
         error ->
-          Logger.error(
-            "Error while listing group mappings for org #{req.org_id}: #{inspect(error)}"
-          )
+          Logger.error("Error while describing mappings for org #{req.org_id}: #{inspect(error)}")
 
-          grpc_error!(:unknown, "Unknown error while listing group mappings")
+          grpc_error!(:unknown, "Unknown error while describing mappings")
       end
     end)
   end

ee/rbac/lib/rbac/okta/idp_group_mapping.ex

@@ -7,31 +7,36 @@ defmodule Rbac.Okta.IdpGroupMapping do
   alias Rbac.Repo.IdpGroupMapping
 
   @doc """
-  Creates or updates group mappings for an organization.
+  Creates or updates mappings for an organization.
 
   ## Parameters
     * organization_id - The ID of the organization
     * group_mappings - A list of maps with idp_group_id and semaphore_group_id keys
-    * default_role_id - The default role ID to use when no group mapping matches
+    * role_mappings - A list of maps with idp_role_id and semaphore_role_id keys
+    * default_role_id - The default role ID to use when no mappings match
 
   ## Returns
     * `{:ok, mapping}` - The created or updated mapping
     * `{:error, changeset}` - Error with changeset containing validation errors
   """
-  def create_or_update(organization_id, group_mapping, default_role_id)
-      when is_list(group_mapping) do
+  def create_or_update(organization_id, group_mapping, role_mapping \\ [], default_role_id)
+      when is_list(group_mapping) and is_list(role_mapping) do
     require Logger
-    Logger.info("req: group_mapping: #{inspect(group_mapping)}")
+
+    Logger.info(
+      "req: group_mapping: #{inspect(group_mapping)}, role_mapping: #{inspect(role_mapping)}"
+    )
 
     IdpGroupMapping.insert_or_update(
       organization_id: organization_id,
       group_mapping: group_mapping,
+      role_mapping: role_mapping,
       default_role_id: default_role_id
     )
   end
 
   @doc """
-  Retrieves group mappings for an organization.
+  Retrieves mappings for an organization.
 
   ## Parameters
     * organization_id - The ID of the organization
@@ -61,7 +66,7 @@ defmodule Rbac.Okta.IdpGroupMapping do
       {:ok, mapping} ->
         # Create a lookup map for faster search
         lookup_map =
-          Enum.reduce(mapping.group_mappings, %{}, fn m, acc ->
+          Enum.reduce(mapping.group_mapping, %{}, fn m, acc ->
             Map.put(acc, m.idp_group_id, m.semaphore_group_id)
           end)
 
@@ -82,4 +87,42 @@ defmodule Rbac.Okta.IdpGroupMapping do
         error
     end
   end
+
+  @doc """
+  Maps IDP roles to Semaphore roles.
+
+  ## Parameters
+    * organization_id - The ID of the organization
+    * idp_roles - List of IDP role identifiers
+
+  ## Returns
+    * `{:ok, semaphore_roles}` - List of mapped Semaphore role IDs
+    * `{:error, :not_found}` - No mapping found for the organization
+  """
+  def map_roles(organization_id, idp_roles) when is_list(idp_roles) do
+    case get_for_organization(organization_id) do
+      {:ok, mapping} ->
+        # Create a lookup map for faster search
+        lookup_map =
+          Enum.reduce(mapping.role_mapping || [], %{}, fn m, acc ->
+            Map.put(acc, m.idp_role_id, m.semaphore_role_id)
+          end)
+
+        # Find matching semaphore roles
+        mapped_roles =
+          idp_roles
+          |> Enum.reduce([], fn idp_role, acc ->
+            case Map.get(lookup_map, idp_role) do
+              nil -> acc
+              semaphore_role -> [semaphore_role | acc]
+            end
+          end)
+          |> Enum.uniq()
+
+        {:ok, mapped_roles}
+
+      error ->
+        error
+    end
+  end
 end