Limit MFA Extension (#2607)

* Implement MFA limit extension --------- Co-authored-by: Stuart Corbishley <[email protected]>
OpenFn · Oct 29, 2024 · 383c692 · 383c692
1 parent a706d75
commit 383c692
Show file tree

Hide file tree

Showing 7 changed files with 145 additions and 18 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -19,6 +19,9 @@ and this project adheres to
 
 ### Changed
 
+- Enforcing MFA for a project can be enforced by the usage limiter
+  [#2607](https://github.com/OpenFn/lightning/pull/2607)
+
 ### Fixed
 
 ## [v2.9.13] - 2024-10-28

diff --git a/lib/lightning/extensions/usage_limiting.ex b/lib/lightning/extensions/usage_limiting.ex
@@ -19,7 +19,8 @@ defmodule Lightning.Extensions.UsageLimiting do
               | :ai_query
               | :alert_failure
               | :github_sync
-              | :new_user,
+              | :new_user
+              | :require_mfa,
             amount: pos_integer()
           }
 
@@ -37,15 +38,55 @@ defmodule Lightning.Extensions.UsageLimiting do
     defstruct [:project_id, :user_id]
   end
 
+  @doc """
+  Checks the usage limits for a given project and user context, returning an error
+  message if the limits are exceeded.
+
+  Requires a `Context` struct containing the `project_id` and `user_id`.
+
+  ## Returns
+
+    - `{:ok}` if the limits are within bounds.
+    - `{:error, reason, message}` if the limits are exceeded.
+  """
   @callback check_limits(context :: Context.t()) ::
               :ok | error()
 
+  @doc """
+  Limits specific actions based on an `Action` and `Context`.
+
+  ## Returns
+
+    - `:ok` if the action is allowed.
+    - `{:error, reason, message}` if the action is not allowed.
+  """
   @callback limit_action(action :: Action.t(), context :: Context.t()) ::
               :ok | error()
 
+  @doc """
+  Increments the AI query count for a given chat session.
+
+  Requires a `Lightning.AiAssistant.ChatSession` struct.
+
+  ## Returns
+
+    - An Ecto.Multi struct representing the transaction to increment AI queries.
+  """
   @callback increment_ai_queries(Lightning.AiAssistant.ChatSession.t()) ::
               Ecto.Multi.t()
 
+  @doc """
+  Returns run options based on the given project context. The run options include
+  a timeout value and a flag for saving dataclips.
+
+  Requires a `Context` struct containing the `project_id`.
+
+  ## Returns
+
+    - A keyword list of run options including:
+      - `:run_timeout_ms`: The run timeout in milliseconds.
+      - `:save_dataclips`: A boolean indicating whether to save dataclips.
+  """
   @callback get_run_options(context :: Context.t()) ::
               Lightning.Runs.RunOptions.keyword_list()
 end
diff --git a/lib/lightning_web/hooks.ex b/lib/lightning_web/hooks.ex
@@ -7,8 +7,11 @@ defmodule LightningWeb.Hooks do
   import Phoenix.Component
   import Phoenix.LiveView
 
+  alias Lightning.Extensions.UsageLimiting.Action
+  alias Lightning.Extensions.UsageLimiting.Context
   alias Lightning.Policies.Permissions
   alias Lightning.Policies.ProjectUsers
+  alias Lightning.Services.UsageLimiter
   alias Lightning.VersionControl.VersionControlUsageLimiter
 
   @doc """
@@ -84,4 +87,19 @@ defmodule LightningWeb.Hooks do
         {:cont, assign(socket, github_banner: component)}
     end
   end
+
+  def on_mount(:limit_mfa, _params, _session, socket) do
+    case UsageLimiter.limit_action(
+           %Action{type: :require_mfa},
+           %Context{
+             project_id: socket.assigns.project.id
+           }
+         ) do
+      :ok ->
+        {:cont, assign(socket, can_require_mfa: true)}
+
+      {:error, _reason, %{function: func} = component} when is_function(func) ->
+        {:cont, assign(socket, mfa_banner: component, can_require_mfa: false)}
+    end
+  end
 end
diff --git a/lib/lightning_web/live/project_live/settings.ex b/lib/lightning_web/live/project_live/settings.ex
@@ -27,6 +27,7 @@ defmodule LightningWeb.ProjectLive.Settings do
 
   on_mount {LightningWeb.Hooks, :project_scope}
   on_mount {LightningWeb.Hooks, :limit_github_sync}
+  on_mount {LightningWeb.Hooks, :limit_mfa}
 
   @impl true
   def mount(_params, _session, socket) do
@@ -194,8 +195,6 @@ defmodule LightningWeb.ProjectLive.Settings do
          ProjectUsers
          |> Permissions.can?(:edit_failure_alerts, current_user, project_user)
 
-  defp can_edit_project(assigns), do: assigns.can_edit_project
-
   @impl true
   def handle_params(params, _url, socket) do
     %{project: %{id: project_id}, live_action: live_action} = socket.assigns
@@ -260,7 +259,7 @@ defmodule LightningWeb.ProjectLive.Settings do
   end
 
   def handle_event("save", %{"project" => project_params}, socket) do
-    if can_edit_project(socket.assigns) do
+    if socket.assigns.can_edit_project do
       save_project(socket, project_params)
     else
       {:noreply,
@@ -284,7 +283,7 @@ defmodule LightningWeb.ProjectLive.Settings do
   end
 
   def handle_event("toggle-mfa", _params, socket) do
-    if can_edit_project(socket.assigns) do
+    if socket.assigns.can_edit_project && socket.assigns.can_require_mfa do
       project = socket.assigns.project
 
       {:ok, project} =

diff --git a/lib/lightning_web/live/project_live/settings.html.heex b/lib/lightning_web/live/project_live/settings.html.heex
@@ -59,7 +59,7 @@
             Projects are isolated workspaces that contain workflows, accessible to certain users.
           </small>
           <.permissions_message
-            :if={!can_edit_project(assigns)}
+            :if={!@can_edit_project}
             section="basic settings, but you can export a copy."
           />
         </div>
@@ -87,7 +87,7 @@
               <div class="md:col-span-2">
                 <Form.text_field
                   form={f}
-                  disabled={!can_edit_project(assigns)}
+                  disabled={!@can_edit_project}
                   label="Project name"
                   field={:name}
                 />
@@ -101,7 +101,7 @@
                 <Form.text_area
                   class="mt-1 focus:ring-primary-500 focus:border-primary-500 block w-full shadow-sm sm:text-sm border-secondary-300 rounded-md"
                   form={f}
-                  disabled={!can_edit_project(assigns)}
+                  disabled={!@can_edit_project}
                   label="Project description"
                   field={:description}
                 />
@@ -116,10 +116,7 @@
             <div class="grid grid-cols gap-6">
               <div class="md:col-span-2">
                 <Form.submit_button
-                  disabled={
-                    !(@project_changeset.valid? and can_edit_project(assigns) and
-                        can_edit_project(assigns))
-                  }
+                  disabled={!(@project_changeset.valid? and @can_edit_project)}
                   phx-disable-with="Saving"
                 >
                   Save
@@ -552,7 +549,7 @@
             View collaborators and manage alert settings for this project.
           </small>
           <.permissions_message
-            :if={!can_edit_project(assigns)}
+            :if={!@can_edit_project}
             section="collaboration settings, but you can change your notification preferences."
           />
         </div>
@@ -681,14 +678,22 @@
             View and manage security settings for this project.
           </small>
           <.permissions_message
-            :if={!can_edit_project(assigns)}
+            :if={!@can_edit_project}
             section="multi-factor authentication settings."
           />
         </div>
         <div class="hidden sm:block" aria-hidden="true">
           <div class="py-2"></div>
         </div>
-
+        <div>
+          <%= if assigns[:mfa_banner] do %>
+            <%= Phoenix.LiveView.TagEngine.component(
+              @mfa_banner.function,
+              @mfa_banner.attrs,
+              {__ENV__.module, __ENV__.function, __ENV__.file, __ENV__.line}
+            ) %>
+          <% end %>
+        </div>
         <div class="bg-white p-4 rounded-md">
           <div class="flex items-center justify-between">
             <span class="flex flex-grow flex-col">
@@ -716,13 +721,14 @@
               <button
                 id="toggle-mfa-switch"
                 type="button"
-                class={"#{if @project.requires_mfa, do: "bg-indigo-600", else: "bg-gray-200"} #{if !can_edit_project(assigns), do: "cursor-not-allowed opacity-50", else: "cursor-pointer"} relative inline-flex h-6 w-11 flex-shrink-0 rounded-full border-2 border-transparent transition-colors duration-200 ease-in-out focus:outline-none focus:ring-2 focus:ring-indigo-600 focus:ring-offset-2"}
+                class={"#{if @project.requires_mfa, do: "bg-indigo-600", else: "bg-gray-200"} #{if !@can_edit_project, do: "cursor-not-allowed opacity-50", else: "cursor-pointer"} relative inline-flex h-6 w-11 flex-shrink-0 rounded-full border-2 border-transparent transition-colors duration-200 ease-in-out focus:outline-none focus:ring-2 focus:ring-indigo-600 focus:ring-offset-2"}
                 role="switch"
+                disabled={!@can_require_mfa}
                 phx-click="toggle-mfa"
                 aria-checked={@project.requires_mfa}
                 aria-labelledby="require-mfa-label"
                 aria-describedby="require-mfa-description"
-                {if !can_edit_project(assigns), do: ["phx-hook": "Tooltip", "data-placement": "bottom", "aria-label": "You do not have permission to perform this action"], else: []}
+                {if !@can_edit_project, do: ["phx-hook": "Tooltip", "data-placement": "bottom", "aria-label": "You do not have permission to perform this action"], else: []}
               >
                 <span
                   aria-hidden="true"

diff --git a/test/lightning_web/live/project_live_test.exs b/test/lightning_web/live/project_live_test.exs
@@ -1451,6 +1451,66 @@ defmodule LightningWeb.ProjectLiveTest do
       assert view |> has_element?("##{form_id}_password_action_button", "Copy")
       assert render(view) =~ auth_method.password
     end
+
+    test "MFA not limited: no banner, button is enabled, and MFA can be toggled",
+         %{conn: conn} do
+      project = insert(:project)
+
+      project_user =
+        insert(:project_user, role: :admin, project: project, user: build(:user))
+
+      stub(Lightning.Extensions.MockUsageLimiter, :limit_action, fn _, _ ->
+        :ok
+      end)
+
+      conn = log_in_user(conn, project_user.user)
+      {:ok, view, _html} = live(conn, ~p"/projects/#{project}/settings#security")
+
+      refute has_element?(view, "#toggle-mfa-switch[disabled]")
+
+      assert view
+             |> element("#toggle-mfa-switch")
+             |> render_click() =~ "Project MFA requirement updated successfully"
+    end
+
+    test "MFA limited: banner displayed, button is disabled, and MFA cannot be toggled",
+         %{conn: conn} do
+      project = insert(:project)
+
+      project_user =
+        insert(:project_user, role: :admin, project: project, user: build(:user))
+
+      banner_message = "MFA Feature is disabled"
+
+      stub(Lightning.Extensions.MockUsageLimiter, :limit_action, fn %{type: type},
+                                                                    _ ->
+        component = %Lightning.Extensions.Message{
+          attrs: %{text: banner_message},
+          function: fn assigns ->
+            ~H"""
+            <%= @text %>
+            """
+          end
+        }
+
+        case type do
+          :require_mfa ->
+            {:error, :disabled, component}
+
+          _ ->
+            :ok
+        end
+      end)
+
+      conn = log_in_user(conn, project_user.user)
+      {:ok, view, html} = live(conn, ~p"/projects/#{project}/settings#security")
+
+      assert html =~ banner_message
+      assert has_element?(view, "#toggle-mfa-switch[disabled]")
+
+      assert view |> render_click("toggle-mfa", %{}) =~
+               "You are not authorized to perform this action"
+    end
   end
 
   test "owners and admins can delete a project webhook auth method",

diff --git a/test/lightning_web/live/workflow_live/edit_test.exs b/test/lightning_web/live/workflow_live/edit_test.exs
@@ -2102,7 +2102,7 @@ defmodule LightningWeb.WorkflowLive.EditTest do
       assert high_priority_view
              |> has_element?("#inspector-workflow-version", "latest")
 
-      retry with: exponential_backoff() |> expiry(4_000),
+      retry with: linear_backoff(50, 2) |> expiry(4_000),
             rescue_only: [ExUnit.AssertionError] do
         refute low_priority_view
                |> has_element?("#inspector-workflow-version", "latest")