Handle authorization groups in real-time (#2998)

Author: aleDsz

lib/livebook/apps.ex

@@ -7,6 +7,7 @@ defmodule Livebook.Apps do
   require Logger
 
   alias Livebook.App
+  alias Livebook.Apps
 
   @doc """
   Returns app process pid for the given slug.
@@ -79,13 +80,10 @@ defmodule Livebook.Apps do
   @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, %{access_type: :full}), do: true
 
-  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)
+  def authorized?(%{slug: slug, app_spec: %Apps.TeamsAppSpec{hub_id: id}}, user) do
+    Livebook.Hubs.TeamClient.user_app_access?(id, user.groups, slug)
   end
 
   @doc """

lib/livebook/hubs/team_client.ex

@@ -296,26 +296,39 @@ 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}
+  def handle_call({:check_full_access, groups}, _caller, state) do
+    if id = state.deployment_group_id do
+      case fetch_deployment_group(id, state) do
+        {:ok, deployment_group} ->
+          {:reply,
+           not deployment_group.teams_auth or
+             not deployment_group.groups_auth or
+             authorized_group?(deployment_group.authorization_groups, groups), state}
+
+        _ ->
+          {:reply, false, state}
+      end
+    else
+      {:reply, true, 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)
+  def handle_call({:check_app_access, groups, slug}, _caller, state) do
+    if id = state.deployment_group_id do
+      with {:ok, deployment_group} <- fetch_deployment_group(id, state),
+           {:ok, app_deployment} <- fetch_app_deployment_from_slug(slug, state) do
+        app_access? =
+          not deployment_group.teams_auth or
+            not deployment_group.groups_auth or
+            (authorized_group?(deployment_group.authorization_groups, groups) or
+               authorized_group?(app_deployment.authorization_groups, groups))
 
-      {:reply, app_access?, state}
+        {:reply, app_access?, state}
+      else
+        _ -> {:reply, false, state}
+      end
     else
-      _ -> {:reply, false, state}
+      {:reply, true, state}
     end
   end
 
@@ -492,6 +505,7 @@ defmodule Livebook.Hubs.TeamClient do
       clustering: nullify(deployment_group.clustering),
       url: nullify(deployment_group.url),
       teams_auth: deployment_group.teams_auth,
+      groups_auth: deployment_group.groups_auth,
       authorization_groups: authorization_groups
     }
   end
@@ -531,6 +545,7 @@ defmodule Livebook.Hubs.TeamClient do
         clustering: atomize(deployment_group_updated.clustering),
         url: nullify(deployment_group_updated.url),
         teams_auth: deployment_group_updated.teams_auth,
+        groups_auth: deployment_group_updated.groups_auth,
         authorization_groups: authorization_groups
     }
   end
@@ -659,7 +674,6 @@ defmodule Livebook.Hubs.TeamClient do
 
   defp handle_event(:deployment_group_created, %Teams.DeploymentGroup{} = deployment_group, state) do
     Teams.Broadcasts.deployment_group_created(deployment_group)
-
     put_deployment_group(state, deployment_group)
   end
 
@@ -673,6 +687,16 @@ defmodule Livebook.Hubs.TeamClient do
 
   defp handle_event(:deployment_group_updated, %Teams.DeploymentGroup{} = deployment_group, state) do
     Teams.Broadcasts.deployment_group_updated(deployment_group)
+
+    with {:ok, current_deployment_group} <- fetch_deployment_group(deployment_group.id, state) do
+      if state.deployment_group_id == deployment_group.id and
+           (current_deployment_group.authorization_groups != deployment_group.authorization_groups or
+              current_deployment_group.groups_auth != deployment_group.groups_auth or
+              current_deployment_group.teams_auth != deployment_group.teams_auth) do
+        Teams.Broadcasts.server_authorization_updated(deployment_group)
+      end
+    end
+
     put_deployment_group(state, deployment_group)
   end
 

lib/livebook/teams/broadcasts.ex

@@ -7,6 +7,7 @@ defmodule Livebook.Teams.Broadcasts do
   @app_deployments_topic "teams:app_deployments"
   @clients_topic "teams:clients"
   @deployment_groups_topic "teams:deployment_groups"
+  @app_server_topic "teams:app_server"
 
   @doc """
   Subscribes to one or more subtopics in `"teams"`.
@@ -31,9 +32,13 @@ defmodule Livebook.Teams.Broadcasts do
   Topic `#{@deployment_groups_topic}`:
 
     * `{:deployment_group_created, DeploymentGroup.t()}`
-    * `{:deployment_group_update, DeploymentGroup.t()}`
+    * `{:deployment_group_updated, DeploymentGroup.t()}`
     * `{:deployment_group_deleted, DeploymentGroup.t()}`
 
+  Topic `#{@app_server_topic}`:
+
+    * `{:server_authorization_updated, DeploymentGroup.t()}`
+
   """
   @spec subscribe(atom() | list(atom())) :: :ok | {:error, term()}
   def subscribe(topics) when is_list(topics) do
@@ -132,6 +137,14 @@ defmodule Livebook.Teams.Broadcasts do
     broadcast(@agents_topic, {:agent_left, agent})
   end
 
+  @doc """
+  Broadcasts under `#{@app_server_topic}` topic when hub received a updated deployment group that changed which groups have access to the server.
+  """
+  @spec server_authorization_updated(Teams.DeploymentGroup.t()) :: broadcast()
+  def server_authorization_updated(%Teams.DeploymentGroup{} = deployment_group) do
+    broadcast(@app_server_topic, {:server_authorization_updated, deployment_group})
+  end
+
   defp broadcast(topic, message) do
     Phoenix.PubSub.broadcast(Livebook.PubSub, topic, message)
   end

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(),
+          groups_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()),
@@ -27,6 +28,7 @@ defmodule Livebook.Teams.DeploymentGroup do
     field :clustering, Ecto.Enum, values: [:auto, :dns]
     field :url, :string
     field :teams_auth, :boolean, default: true
+    field :groups_auth, :boolean, default: false
 
     has_many :secrets, Secrets.Secret
     has_many :agent_keys, Teams.AgentKey
@@ -36,7 +38,7 @@ defmodule Livebook.Teams.DeploymentGroup do
 
   def changeset(deployment_group, attrs \\ %{}) do
     deployment_group
-    |> cast(attrs, [:id, :name, :mode, :hub_id, :clustering, :url, :teams_auth])
+    |> cast(attrs, [:id, :name, :mode, :hub_id, :clustering, :url])
     |> validate_required([:name, :mode])
     |> update_change(:url, fn url ->
       if url do

lib/livebook/teams/requests.ex

@@ -172,8 +172,7 @@ defmodule Livebook.Teams.Requests do
       name: deployment_group.name,
       mode: deployment_group.mode,
       clustering: deployment_group.clustering,
-      url: deployment_group.url,
-      teams_auth: deployment_group.teams_auth
+      url: deployment_group.url
     }
 
     post("/api/v1/org/deployment-groups", params, team)

lib/livebook/users/user.ex

@@ -12,12 +12,15 @@ defmodule Livebook.Users.User do
 
   alias Livebook.Utils
 
+  @type access_type :: :full | :apps
+
   @type t :: %__MODULE__{
           id: id(),
           name: String.t() | nil,
           email: String.t() | nil,
           avatar_url: String.t() | nil,
-          restricted_apps_groups: list(map()) | nil,
+          access_type: access_type(),
+          groups: list(map()) | nil,
           payload: map() | nil,
           hex_color: hex_color()
         }
@@ -29,7 +32,8 @@ defmodule Livebook.Users.User do
     field :name, :string
     field :email, :string
     field :avatar_url, :string
-    field :restricted_apps_groups, {:array, :map}
+    field :access_type, Ecto.Enum, values: ~w[full apps]a, default: :full
+    field :groups, {:array, :map}, default: []
     field :payload, :map
     field :hex_color, Livebook.EctoTypes.HexColor
   end
@@ -44,15 +48,17 @@ defmodule Livebook.Users.User do
       name: nil,
       email: nil,
       avatar_url: nil,
-      restricted_apps_groups: nil,
+      access_type: :full,
+      groups: [],
       payload: nil,
       hex_color: Livebook.EctoTypes.HexColor.random()
     }
   end
 
+  @doc false
   def changeset(user, attrs \\ %{}) do
     user
-    |> cast(attrs, [:name, :email, :avatar_url, :restricted_apps_groups, :hex_color, :payload])
+    |> cast(attrs, [:name, :email, :avatar_url, :access_type, :groups, :hex_color, :payload])
     |> validate_required([:hex_color])
   end
 end

lib/livebook/zta.ex

@@ -58,6 +58,9 @@ defmodule Livebook.ZTA do
     * `:id` - a string that uniquely identifies the user
     * `:name` - the user name
     * `:email` - the user email
+    * `:avatar_url` - the user avatar
+    * `:access_type` - the user access type
+    * `:groups` - the user groups
     * `:payload` - the provider payload
 
   Note that none of the keys are required. The metadata returned depends
@@ -67,6 +70,9 @@ defmodule Livebook.ZTA do
           optional(:id) => String.t(),
           optional(:name) => String.t(),
           optional(:email) => String.t(),
+          optional(:avatar_url) => String.t() | nil,
+          optional(:access_type) => Livebook.Users.User.access_type(),
+          optional(:groups) => list(map()),
           optional(:payload) => map()
         }
 

lib/livebook/zta/livebook_teams.ex

@@ -159,31 +159,36 @@ defmodule Livebook.ZTA.LivebookTeams do
 
   defp get_user_info(team, access_token) do
     with {:ok, payload} <- Teams.Requests.get_user_info(team, access_token) do
-      %{
-        "id" => id,
-        "name" => name,
-        "email" => email,
-        "groups" => groups,
-        "avatar_url" => avatar_url
-      } = payload
-
-      restricted_apps_groups =
-        if Livebook.Hubs.TeamClient.user_full_access?(team.id, groups) do
-          nil
-        else
-          groups
-        end
-
-      metadata = %{
-        id: id,
-        name: name,
-        avatar_url: avatar_url,
-        restricted_apps_groups: restricted_apps_groups,
-        email: email,
-        payload: payload
-      }
-
-      {:ok, metadata}
+      {:ok, build_metadata(team.id, payload)}
     end
   end
+
+  @doc """
+  Returns the user metadata from given payload.
+  """
+  @spec build_metadata(String.t(), map()) :: Livebook.ZTA.metadata()
+  def build_metadata(hub_id, payload) do
+    %{
+      "id" => id,
+      "name" => name,
+      "email" => email,
+      "groups" => groups,
+      "avatar_url" => avatar_url
+    } = payload
+
+    access_type =
+      if Livebook.Hubs.TeamClient.user_full_access?(hub_id, groups),
+        do: :full,
+        else: :apps
+
+    %{
+      id: id,
+      name: name,
+      avatar_url: avatar_url,
+      access_type: access_type,
+      groups: groups,
+      email: email,
+      payload: payload
+    }
+  end
 end

lib/livebook_web/live/app_live.ex

@@ -12,8 +12,13 @@ defmodule LivebookWeb.AppLive do
     end
   end
 
-  def mount(_params, _session, socket) when not socket.assigns.app_authorized? do
-    {:ok, socket, layout: false}
+  def mount(%{"slug" => slug}, _session, socket) when not socket.assigns.app_authorized? do
+    if connected?(socket) do
+      Livebook.Teams.Broadcasts.subscribe(:app_deployments)
+    end
+
+    {:ok, app} = Livebook.Apps.fetch_app(slug)
+    {:ok, assign(socket, app: app), layout: false}
   end
 
   def mount(%{"slug" => slug}, _session, socket) do
@@ -22,6 +27,7 @@ defmodule LivebookWeb.AppLive do
 
       if connected?(socket) do
         Livebook.App.subscribe(slug)
+        Livebook.Teams.Broadcasts.subscribe(:app_deployments)
       end
 
       {:ok, assign(socket, app: app)}
@@ -122,6 +128,18 @@ defmodule LivebookWeb.AppLive do
     {:noreply, assign(socket, :app, app)}
   end
 
+  def handle_info(
+        {:app_deployment_updated, %{slug: slug}},
+        %{assigns: %{app: %{slug: slug} = app}} = socket
+      ) do
+    if socket.assigns.app_authorized? and
+         Livebook.Apps.authorized?(app, socket.assigns.current_user) do
+      {:noreply, socket}
+    else
+      {:noreply, redirect(socket, to: ~p"/apps/#{slug}")}
+    end
+  end
+
   def handle_info(_message, socket), do: {:noreply, socket}
 
   defp active_sessions(sessions) do

lib/livebook_web/live/app_session_live.ex

@@ -24,8 +24,13 @@ defmodule LivebookWeb.AppSessionLive do
     end
   end
 
-  def mount(_params, _session, socket) when not socket.assigns.app_authorized? do
-    {:ok, socket, layout: false}
+  def mount(%{"slug" => slug, "id" => session_id}, _session, socket)
+      when not socket.assigns.app_authorized? do
+    if connected?(socket) do
+      Livebook.Teams.Broadcasts.subscribe(:app_deployments)
+    end
+
+    {:ok, assign(socket, slug: slug, session: %{id: session_id}), layout: false}
   end
 
   def mount(%{"slug" => slug, "id" => session_id}, _session, socket) do
@@ -388,7 +393,8 @@ defmodule LivebookWeb.AppSessionLive do
     # should't have access.
     {:ok, app} = Livebook.Apps.fetch_app(slug)
 
-    if Livebook.Apps.authorized?(app, socket.assigns.current_user) do
+    if socket.assigns.app_authorized? and
+         Livebook.Apps.authorized?(app, socket.assigns.current_user) do
       {:noreply, socket}
     else
       {:noreply, redirect(socket, to: ~p"/apps/#{slug}/sessions/#{socket.assigns.session.id}")}