Implement the user authorization flow with Livebook Teams

lib/livebook/apps.ex

@@ -62,6 +62,32 @@ defmodule Livebook.Apps do
     Livebook.Tracker.list_apps()
   end
 
+  @doc """
+  Returns all the running apps authorized to given user.
+  """
+  @spec list_authorized_apps(Livebook.Users.User.t()) :: list(App.t())
+  def list_authorized_apps(user) do
+    for app <- list_apps(),
+        authorized?(app, user) do
+      app
+    end
+  end
+
+  @doc """
+  Returns if the given running app is authorized to given user.
+  """
+  @spec authorized?(App.t(), Livebook.Users.User.t()) :: boolean()
+  def authorized?(app, user)
+
+  def authorized?(%{app_spec: %Livebook.Apps.TeamsAppSpec{}}, %{restricted_apps_groups: []}),
+    do: false
+
+  def authorized?(_app, %{restricted_apps_groups: nil}), do: true
+
+  def authorized?(%{slug: slug, app_spec: %Livebook.Apps.TeamsAppSpec{hub_id: id}}, user) do
+    Livebook.Hubs.TeamClient.user_app_access?(id, user.restricted_apps_groups, slug)
+  end
+
   @doc """
   Updates the given app info across the cluster.
   """

lib/livebook/config.ex

@@ -6,6 +6,8 @@ defmodule Livebook.Config do
           | %{mode: :token, secret: String.t()}
           | %{mode: :disabled}
 
+  @type authentication_mode :: :password | :token | :disabled
+
   @doc """
   Returns path to Livebook priv directory.
 

lib/livebook/hubs/team_client.ex

@@ -142,6 +142,22 @@ defmodule Livebook.Hubs.TeamClient do
     GenServer.call(registry_name(id), :get_environment_variables)
   end
 
+  @doc """
+  Returns if the given user groups has full access to app server.
+  """
+  @spec user_full_access?(String.t(), list(map())) :: boolean()
+  def user_full_access?(id, groups) do
+    GenServer.call(registry_name(id), {:check_full_access, groups})
+  end
+
+  @doc """
+  Returns if the given user groups has access to given app.
+  """
+  @spec user_app_access?(String.t(), list(map()), String.t()) :: boolean()
+  def user_app_access?(id, groups, slug) do
+    GenServer.call(registry_name(id), {:check_app_access, groups, slug})
+  end
+
   @doc """
   Returns if the Team client is connected.
   """
@@ -280,6 +296,29 @@ defmodule Livebook.Hubs.TeamClient do
     end
   end
 
+  def handle_call({:check_full_access, groups}, _caller, %{deployment_group_id: id} = state) do
+    case fetch_deployment_group(id, state) do
+      {:ok, deployment_group} ->
+        {:reply, authorized_group?(deployment_group.authorization_groups, groups), state}
+
+      _ ->
+        {:reply, false, state}
+    end
+  end
+
+  def handle_call({:check_app_access, groups, slug}, _caller, %{deployment_group_id: id} = state) do
+    with {:ok, deployment_group} <- fetch_deployment_group(id, state),
+         {:ok, app_deployment} <- fetch_app_deployment_from_slug(slug, state) do
+      app_access? =
+        authorized_group?(deployment_group.authorization_groups, groups) or
+          authorized_group?(app_deployment.authorization_groups, groups)
+
+      {:reply, app_access?, state}
+    else
+      _ -> {:reply, false, state}
+    end
+  end
+
   @impl true
   def handle_info(:connected, state) do
     Hubs.Broadcasts.hub_connected(state.hub.id)
@@ -440,6 +479,7 @@ defmodule Livebook.Hubs.TeamClient do
     secrets = Enum.map(deployment_group.secrets, &build_secret(state, &1))
     agent_keys = Enum.map(deployment_group.agent_keys, &build_agent_key/1)
     environment_variables = build_environment_variables(state, deployment_group)
+    authorization_groups = build_authorization_groups(deployment_group)
 
     %Teams.DeploymentGroup{
       id: deployment_group.id,
@@ -451,7 +491,8 @@ defmodule Livebook.Hubs.TeamClient do
       environment_variables: environment_variables,
       clustering: nullify(deployment_group.clustering),
       url: nullify(deployment_group.url),
-      teams_auth: deployment_group.teams_auth
+      teams_auth: deployment_group.teams_auth,
+      authorization_groups: authorization_groups
     }
   end
 
@@ -468,14 +509,16 @@ defmodule Livebook.Hubs.TeamClient do
       environment_variables: [],
       clustering: nullify(deployment_group_created.clustering),
       url: nullify(deployment_group_created.url),
-      teams_auth: deployment_group_created.teams_auth
+      teams_auth: deployment_group_created.teams_auth,
+      authorization_groups: []
     }
   end
 
   defp build_deployment_group(state, deployment_group_updated) do
     secrets = Enum.map(deployment_group_updated.secrets, &build_secret(state, &1))
     agent_keys = Enum.map(deployment_group_updated.agent_keys, &build_agent_key/1)
     environment_variables = build_environment_variables(state, deployment_group_updated)
+    authorization_groups = build_authorization_groups(deployment_group_updated)
 
     {:ok, deployment_group} = fetch_deployment_group(deployment_group_updated.id, state)
 
@@ -487,11 +530,14 @@ defmodule Livebook.Hubs.TeamClient do
         environment_variables: environment_variables,
         clustering: atomize(deployment_group_updated.clustering),
         url: nullify(deployment_group_updated.url),
-        teams_auth: deployment_group_updated.teams_auth
+        teams_auth: deployment_group_updated.teams_auth,
+        authorization_groups: authorization_groups
     }
   end
 
   defp build_app_deployment(state, %LivebookProto.AppDeployment{} = app_deployment) do
+    authorization_groups = build_authorization_groups(app_deployment)
+
     %Teams.AppDeployment{
       id: app_deployment.id,
       slug: app_deployment.slug,
@@ -504,7 +550,8 @@ defmodule Livebook.Hubs.TeamClient do
       deployment_group_id: app_deployment.deployment_group_id,
       file: nil,
       deployed_by: app_deployment.deployed_by,
-      deployed_at: DateTime.from_gregorian_seconds(app_deployment.deployed_at)
+      deployed_at: DateTime.from_gregorian_seconds(app_deployment.deployed_at),
+      authorization_groups: authorization_groups
     }
   end
 
@@ -519,6 +566,15 @@ defmodule Livebook.Hubs.TeamClient do
     end
   end
 
+  defp build_authorization_groups(%{authorization_groups: authorization_groups}) do
+    for authorization_group <- authorization_groups do
+      %Teams.AuthorizationGroup{
+        provider_id: authorization_group.provider_id,
+        group_name: authorization_group.group_name
+      }
+    end
+  end
+
   defp put_agent(state, agent) do
     state = remove_agent(state, agent)
 
@@ -662,14 +718,7 @@ defmodule Livebook.Hubs.TeamClient do
   end
 
   defp handle_event(:app_deployment_started, %Teams.AppDeployment{} = app_deployment, state) do
-    deployment_group_id = app_deployment.deployment_group_id
-
-    with {:ok, deployment_group} <- fetch_deployment_group(deployment_group_id, state) do
-      if deployment_group.id == state.deployment_group_id do
-        manager_sync()
-      end
-    end
-
+    manager_sync(app_deployment, state)
     Teams.Broadcasts.app_deployment_started(app_deployment)
     put_app_deployment(state, app_deployment)
   end
@@ -683,14 +732,7 @@ defmodule Livebook.Hubs.TeamClient do
   end
 
   defp handle_event(:app_deployment_stopped, %Teams.AppDeployment{} = app_deployment, state) do
-    deployment_group_id = app_deployment.deployment_group_id
-
-    with {:ok, deployment_group} <- fetch_deployment_group(deployment_group_id, state) do
-      if deployment_group.id == state.deployment_group_id do
-        manager_sync()
-      end
-    end
-
+    manager_sync(app_deployment, state)
     Teams.Broadcasts.app_deployment_stopped(app_deployment)
     remove_app_deployment(state, app_deployment)
   end
@@ -701,6 +743,20 @@ defmodule Livebook.Hubs.TeamClient do
     end
   end
 
+  defp handle_event(:app_deployment_updated, %Teams.AppDeployment{} = app_deployment, state) do
+    manager_sync(app_deployment, state)
+    Teams.Broadcasts.app_deployment_updated(app_deployment)
+    put_app_deployment(state, app_deployment)
+  end
+
+  defp handle_event(:app_deployment_updated, app_deployment_updated, state) do
+    handle_event(
+      :app_deployment_updated,
+      build_app_deployment(state, app_deployment_updated.app_deployment),
+      state
+    )
+  end
+
   defp handle_event(:agent_joined, %Teams.Agent{} = agent, state) do
     Teams.Broadcasts.agent_joined(agent)
     put_agent(state, agent)
@@ -921,6 +977,9 @@ defmodule Livebook.Hubs.TeamClient do
   defp fetch_app_deployment(id, state),
     do: fetch_entry(state.app_deployments, &(&1.id == id), state)
 
+  defp fetch_app_deployment_from_slug(slug, state),
+    do: fetch_entry(state.app_deployments, &(&1.slug == slug), state)
+
   defp fetch_entry(entries, fun, state) do
     if entry = Enum.find(entries, fun) do
       {:ok, entry}
@@ -938,10 +997,20 @@ defmodule Livebook.Hubs.TeamClient do
   defp nullify(""), do: nil
   defp nullify(value), do: value
 
-  defp manager_sync() do
-    # Each node runs the teams client, but we only need to call sync once
-    if Apps.Manager.local?() do
-      Apps.Manager.sync_permanent_apps()
+  defp manager_sync(app_deployment, state) do
+    # We only need to sync if the app deployment belongs to the current
+    # deployment group
+    if app_deployment.deployment_group_id == state.deployment_group_id do
+      # Each node runs the teams client, but we only need to call sync once
+      if Apps.Manager.local?() do
+        Apps.Manager.sync_permanent_apps()
+      end
     end
   end
+
+  defp authorized_group?(authorization_groups, groups) do
+    Enum.any?(authorization_groups, fn %{provider_id: id, group_name: name} ->
+      %{"provider_id" => id, "group_name" => name} in groups
+    end)
+  end
 end

lib/livebook/teams/app_deployment.ex

@@ -16,7 +16,8 @@ defmodule Livebook.Teams.AppDeployment do
           deployment_group_id: String.t() | nil,
           file: binary() | nil,
           deployed_by: String.t() | nil,
-          deployed_at: DateTime.t() | nil
+          deployed_at: DateTime.t() | nil,
+          authorization_groups: Ecto.Schema.embeds_many(Livebook.Teams.AuthorizationGroup.t())
         }
 
   @access_types Livebook.Notebook.AppSettings.access_types()
@@ -34,6 +35,8 @@ defmodule Livebook.Teams.AppDeployment do
     field :file, :string
     field :deployed_by, :string
     field :deployed_at, :utc_datetime
+
+    embeds_many :authorization_groups, Livebook.Teams.AuthorizationGroup
   end
 
   @doc """

lib/livebook/teams/authorization_group.ex

@@ -0,0 +1,14 @@
+defmodule Livebook.Teams.AuthorizationGroup do
+  use Ecto.Schema
+
+  @type t :: %__MODULE__{
+          provider_id: String.t() | nil,
+          group_name: String.t() | nil
+        }
+
+  @primary_key false
+  embedded_schema do
+    field :provider_id, :string
+    field :group_name, :string
+  end
+end

lib/livebook/teams/broadcasts.ex

@@ -22,6 +22,7 @@ defmodule Livebook.Teams.Broadcasts do
 
     * `{:app_deployment_started, AppDeployment.t()}`
     * `{:app_deployment_stopped, AppDeployment.t()}`
+    * `{:app_deployment_updated, AppDeployment.t()}`
 
   Topic `#{@clients_topic}`:
 
@@ -107,6 +108,14 @@ defmodule Livebook.Teams.Broadcasts do
     broadcast(@app_deployments_topic, {:app_deployment_stopped, app_deployment})
   end
 
+  @doc """
+  Broadcasts under `#{@app_deployments_topic}` topic when hub received an updated app deployment.
+  """
+  @spec app_deployment_updated(Teams.AppDeployment.t()) :: broadcast()
+  def app_deployment_updated(%Teams.AppDeployment{} = app_deployment) do
+    broadcast(@app_deployments_topic, {:app_deployment_updated, app_deployment})
+  end
+
   @doc """
   Broadcasts under `#{@agents_topic}` topic when hub received a new agent.
   """

lib/livebook/teams/deployment_group.ex

@@ -13,6 +13,7 @@ defmodule Livebook.Teams.DeploymentGroup do
           clustering: :auto | :dns | nil,
           hub_id: String.t() | nil,
           teams_auth: boolean(),
+          authorization_groups: Ecto.Schema.embeds_many(Teams.AuthorizationGroup.t()),
           secrets: Ecto.Schema.has_many(Secrets.Secret.t()),
           agent_keys: Ecto.Schema.has_many(Teams.AgentKey.t()),
           environment_variables: Ecto.Schema.has_many(Teams.EnvironmentVariable.t())
@@ -30,6 +31,7 @@ defmodule Livebook.Teams.DeploymentGroup do
     has_many :secrets, Secrets.Secret
     has_many :agent_keys, Teams.AgentKey
     has_many :environment_variables, Teams.EnvironmentVariable
+    embeds_many :authorization_groups, Teams.AuthorizationGroup
   end
 
   def changeset(deployment_group, attrs \\ %{}) do

lib/livebook/users/user.ex

@@ -17,6 +17,7 @@ defmodule Livebook.Users.User do
           name: String.t() | nil,
           email: String.t() | nil,
           avatar_url: String.t() | nil,
+          restricted_apps_groups: list(map()) | nil,
           payload: map() | nil,
           hex_color: hex_color()
         }
@@ -28,6 +29,7 @@ defmodule Livebook.Users.User do
     field :name, :string
     field :email, :string
     field :avatar_url, :string
+    field :restricted_apps_groups, {:array, :map}
     field :payload, :map
     field :hex_color, Livebook.EctoTypes.HexColor
   end
@@ -42,14 +44,15 @@ defmodule Livebook.Users.User do
       name: nil,
       email: nil,
       avatar_url: nil,
+      restricted_apps_groups: nil,
       payload: nil,
       hex_color: Livebook.EctoTypes.HexColor.random()
     }
   end
 
   def changeset(user, attrs \\ %{}) do
     user
-    |> cast(attrs, [:name, :email, :avatar_url, :hex_color, :payload])
+    |> cast(attrs, [:name, :email, :avatar_url, :restricted_apps_groups, :hex_color, :payload])
     |> validate_required([:hex_color])
   end
 end