Add support for reading app info from the runtime (#2038)

README.md

@@ -209,7 +209,7 @@ The following environment variables can be used to configure Livebook on boot:
     Supported values are:
 
       * "cloudflare:<your-team-name (domain)>"
-      * "googleiap:<your-audience (aud)>"
+      * "google_iap:<your-audience (aud)>"
 
   * LIVEBOOK_IFRAME_PORT - sets the port that Livebook serves iframes at.
     This is relevant only when running Livebook without TLS. Defaults to 8081.

lib/livebook/app.ex

@@ -215,14 +215,6 @@ defmodule Livebook.App do
   def handle_info({:DOWN, _ref, :process, pid, _reason}, state) do
     app_session = Enum.find(state.sessions, &(&1.pid == pid))
     state = update_in(state.sessions, &(&1 -- [app_session]))
-
-    state =
-      if user_id = app_session.started_by_id do
-        untrack_user(state, user_id)
-      else
-        state
-      end
-
     {:noreply, notify_update(state)}
   end
 
@@ -240,12 +232,7 @@ defmodule Livebook.App do
       notebook_name: state.notebook.name,
       public?: state.notebook.app_settings.access_type == :public,
       multi_session: state.notebook.app_settings.multi_session,
-      sessions:
-        for session <- state.sessions do
-          {started_by_id, session} = Map.pop!(session, :started_by_id)
-          started_by = started_by_id && state.users[started_by_id].user
-          Map.put(session, :started_by, started_by)
-        end
+      sessions: state.sessions
     }
   end
 
@@ -274,11 +261,14 @@ defmodule Livebook.App do
   end
 
   defp start_app_session(state, user \\ nil) do
+    user = if(state.notebook.teams_enabled, do: user)
+
     opts = [
       notebook: state.notebook,
       mode: :app,
       app_pid: self(),
-      auto_shutdown_ms: state.notebook.app_settings.auto_shutdown_ms
+      auto_shutdown_ms: state.notebook.app_settings.auto_shutdown_ms,
+      started_by: user
     ]
 
     case Livebook.Sessions.create_session(opts) do
@@ -290,20 +280,13 @@ defmodule Livebook.App do
           created_at: session.created_at,
           app_status: %{execution: :executing, lifecycle: :active},
           client_count: 0,
-          started_by_id: user && user.id
+          started_by: user
         }
 
         Process.monitor(session.pid)
 
         state = update_in(state.sessions, &[app_session | &1])
 
-        state =
-          if user do
-            track_user(state, user)
-          else
-            state
-          end
-
         {:ok, state, app_session}
 
       {:error, reason} ->
@@ -352,23 +335,4 @@ defmodule Livebook.App do
   defp broadcast_message(slug, message) do
     Phoenix.PubSub.broadcast(Livebook.PubSub, "apps:#{slug}", message)
   end
-
-  defp track_user(state, user) do
-    if Map.has_key?(state.users, user.id) do
-      update_in(state.users[user.id].count, &(&1 + 1))
-    else
-      Livebook.Users.subscribe(user.id)
-      put_in(state.users[user.id], %{user: user, count: 1})
-    end
-  end
-
-  defp untrack_user(state, user_id) do
-    if state.users[user_id] == 1 do
-      {_, state} = pop_in(state.users[user_id])
-      Livebook.Users.unsubscribe(user_id)
-      state
-    else
-      update_in(state.users[user_id].count, &(&1 - 1))
-    end
-  end
 end

lib/livebook/config.ex

@@ -5,6 +5,20 @@ defmodule Livebook.Config do
 
   @type auth_mode() :: :token | :password | :disabled
 
+  identity_providers = %{
+    session: LivebookWeb.SessionIdentity,
+    google_iap: Livebook.ZTA.GoogleIAP,
+    cloudflare: Livebook.ZTA.Cloudflare
+  }
+
+  @identity_provider_type_to_module Map.new(identity_providers, fn {key, value} ->
+                                      {Atom.to_string(key), value}
+                                    end)
+
+  @identity_provider_module_to_type Map.new(identity_providers, fn {key, value} ->
+                                      {value, key}
+                                    end)
+
   @doc """
   Returns the longname if the distribution mode is configured to use long names.
   """
@@ -184,6 +198,15 @@ defmodule Livebook.Config do
     not match?({LivebookWeb.SessionIdentity, _}, Livebook.Config.identity_provider())
   end
 
+  @doc """
+  Returns identity source as a friendly atom.
+  """
+  @spec identity_source() :: atom()
+  def identity_source() do
+    {module, _} = identity_provider()
+    @identity_provider_module_to_type[module]
+  end
+
   @doc """
   Returns whether the application is running inside an iframe.
   """
@@ -514,18 +537,20 @@ defmodule Livebook.Config do
   Parses zero trust identity provider from env.
   """
   def identity_provider!(env) do
-    case System.get_env(env) do
-      "googleiap:" <> rest ->
-        {Livebook.ZTA.GoogleIAP, rest}
+    if provider = System.get_env(env) do
+      identity_provider!(env, provider)
+    end
+  end
 
-      "cloudflare:" <> rest ->
-        {Livebook.ZTA.Cloudflare, rest}
-
-      nil ->
-        nil
-
-      _ ->
-        abort!("invalid configuration for identity provider")
+  @doc """
+  Parses and validates zero trust identity provider within context.
+  """
+  def identity_provider!(context, provider) do
+    with [type, key] <- String.split(provider, ":", parts: 2),
+         {:ok, module} <- Map.fetch(@identity_provider_type_to_module, type) do
+      {module, key}
+    else
+      _ -> abort!("invalid configuration for identity provider given in #{context}")
     end
   end
 end

lib/livebook/live_markdown/import.ex

@@ -527,16 +527,28 @@ defmodule Livebook.LiveMarkdown.Import do
   defp postprocess_stamp(notebook, _notebook_source, nil, _), do: {notebook, []}
 
   defp postprocess_stamp(notebook, notebook_source, stamp_data, stamp_hub_id) do
-    hub = Hubs.get_offline_hub(stamp_hub_id) || Hubs.fetch_hub!(stamp_hub_id)
+    {hub, offline?} =
+      cond do
+        hub = Hubs.get_offline_hub(stamp_hub_id) -> {hub, true}
+        hub = Hubs.fetch_hub!(stamp_hub_id) -> {hub, false}
+      end
 
-    with %{"offset" => offset, "stamp" => stamp} <- stamp_data,
-         {:ok, notebook_source} <- safe_binary_slice(notebook_source, 0, offset),
-         {:ok, metadata} <- Livebook.Hubs.verify_notebook_stamp(hub, notebook_source, stamp) do
-      notebook = apply_stamp_metadata(notebook, metadata)
-      {notebook, []}
-    else
-      _ -> {notebook, ["failed to verify notebook stamp"]}
-    end
+    {valid_stamp?, notebook, messages} =
+      with %{"offset" => offset, "stamp" => stamp} <- stamp_data,
+           {:ok, notebook_source} <- safe_binary_slice(notebook_source, 0, offset),
+           {:ok, metadata} <- Livebook.Hubs.verify_notebook_stamp(hub, notebook_source, stamp) do
+        notebook = apply_stamp_metadata(notebook, metadata)
+        {true, notebook, []}
+      else
+        _ -> {false, notebook, ["failed to verify notebook stamp"]}
+      end
+
+    # We enable teams features for offline hub only if the stamp
+    # is valid, which ensures it is an existing Teams hub
+    teams_enabled = is_struct(hub, Livebook.Hubs.Team) and (not offline? or valid_stamp?)
+
+    notebook = %{notebook | teams_enabled: teams_enabled}
+    {notebook, messages}
   end
 
   defp safe_binary_slice(binary, start, size)

lib/livebook/notebook.ex

@@ -25,7 +25,8 @@ defmodule Livebook.Notebook do
     :output_counter,
     :app_settings,
     :hub_id,
-    :hub_secret_names
+    :hub_secret_names,
+    :teams_enabled
   ]
 
   alias Livebook.Notebook.{Section, Cell, AppSettings}
@@ -44,7 +45,8 @@ defmodule Livebook.Notebook do
           output_counter: non_neg_integer(),
           app_settings: AppSettings.t(),
           hub_id: String.t(),
-          hub_secret_names: list(String.t())
+          hub_secret_names: list(String.t()),
+          teams_enabled: boolean()
         }
 
   @version "1.0"
@@ -66,7 +68,8 @@ defmodule Livebook.Notebook do
       output_counter: 0,
       app_settings: AppSettings.new(),
       hub_id: Livebook.Hubs.Personal.id(),
-      hub_secret_names: []
+      hub_secret_names: [],
+      teams_enabled: false
     }
     |> put_setup_cell(Cell.new(:code))
   end

lib/livebook/runtime/evaluator/io_proxy.ex

@@ -317,6 +317,11 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
     {state.file, state}
   end
 
+  defp io_request(:livebook_get_app_info, state) do
+    result = request_app_info(state)
+    {result, state}
+  end
+
   defp io_request(_, state) do
     {{:error, :request}, state}
   end
@@ -383,6 +388,21 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
     end
   end
 
+  defp request_app_info(state) do
+    send(state.send_to, {:runtime_app_info_request, self()})
+
+    ref = Process.monitor(state.send_to)
+
+    receive do
+      {:runtime_app_info_reply, app_info} ->
+        Process.demonitor(ref, [:flush])
+        {:ok, app_info}
+
+      {:DOWN, ^ref, :process, _object, _reason} ->
+        {:error, :terminated}
+    end
+  end
+
   defp io_reply(from, reply_as, reply) do
     send(from, {:io_reply, reply_as, reply})
   end

lib/livebook/session.ex

@@ -151,6 +151,10 @@ defmodule Livebook.Session do
     * `:auto_shutdown_ms` - the inactivity period (no clients) after which
       the session should close automatically
 
+    * `:started_by` - the user that started the session. This is relevant
+      for app sessions using the Teams hub, in which case this information
+      is accessible from runtime
+
   """
   @spec start_link(keyword()) :: {:ok, pid} | {:error, any()}
   def start_link(opts) do
@@ -746,7 +750,8 @@ defmodule Livebook.Session do
         deployed_app_monitor_ref: nil,
         app_pid: opts[:app_pid],
         auto_shutdown_ms: opts[:auto_shutdown_ms],
-        auto_shutdown_timer_ref: nil
+        auto_shutdown_timer_ref: nil,
+        started_by: opts[:started_by]
       }
 
       {:ok, state}
@@ -1392,6 +1397,11 @@ defmodule Livebook.Session do
     {:noreply, state}
   end
 
+  def handle_info({:runtime_app_info_request, reply_to}, state) do
+    send(reply_to, {:runtime_app_info_reply, app_info_for_runtime(state)})
+    {:noreply, state}
+  end
+
   def handle_info({:runtime_container_down, container_ref, message}, state) do
     broadcast_error(state.session_id, "evaluation process terminated - #{message}")
 
@@ -2268,6 +2278,30 @@ defmodule Livebook.Session do
     broadcast_message(state.session_id, :session_closed)
   end
 
+  defp app_info_for_runtime(state) do
+    case state.data do
+      %{mode: :app, notebook: %{app_settings: %{multi_session: true}}} ->
+        info = %{type: :multi_session}
+
+        if user = state.started_by do
+          started_by =
+            user
+            |> Map.take([:id, :name, :email])
+            |> Map.put(:source, Livebook.Config.identity_source())
+
+          Map.put(info, :started_by, started_by)
+        else
+          info
+        end
+
+      %{mode: :app, notebook: %{app_settings: %{multi_session: false}}} ->
+        %{type: :single_session}
+
+      _ ->
+        %{type: :none}
+    end
+  end
+
   defp app_report_client_count_change(state) when state.data.mode == :app do
     client_count = map_size(state.data.clients_map)
     send(state.app_pid, {:app_client_count_changed, state.session_id, client_count})

lib/livebook/session/data.ex

@@ -1637,7 +1637,11 @@ defmodule Livebook.Session.Data do
   defp set_notebook_hub({data, _} = data_actions, hub) do
     data_actions
     |> set!(
-      notebook: %{data.notebook | hub_id: hub.id},
+      notebook: %{
+        data.notebook
+        | hub_id: hub.id,
+          teams_enabled: is_struct(hub, Livebook.Hubs.Team)
+      },
       hub_secrets: Hubs.get_secrets(hub)
     )
   end

lib/livebook/stamping.ex

@@ -52,14 +52,20 @@ defmodule Livebook.Stamping do
   @spec rsa_verify?(String.t(), String.t(), String.t()) :: boolean()
   def rsa_verify?(signature, payload, public_key) do
     der_key = Base.url_decode64!(public_key, padding: false)
-    raw_key = :public_key.der_decode(:RSAPublicKey, der_key)
 
-    case Base.url_decode64(signature, padding: false) do
-      {:ok, raw_signature} ->
-        :public_key.verify(payload, :sha256, raw_signature, raw_key)
+    with {:ok, raw_key} <- safe_der_decode(der_key),
+         {:ok, raw_signature} <- Base.url_decode64(signature, padding: false) do
+      :public_key.verify(payload, :sha256, raw_signature, raw_key)
+    else
+      _ -> false
+    end
+  end
 
-      :error ->
-        false
+  defp safe_der_decode(der_key) do
+    try do
+      {:ok, :public_key.der_decode(:RSAPublicKey, der_key)}
+    rescue
+      _ -> :error
     end
   end
 end

lib/livebook/zta/google_iap.ex

@@ -1,7 +1,7 @@
 defmodule Livebook.ZTA.GoogleIAP do
   @doc """
   To integrate your Google Identity-Aware Proxy (IAP) authentication with Livebook,
-  set the `LIVEBOOK_IDENTITY_PROVIDER` environment variable to `googleiap:<your-jwt-audience>`
+  set the `LIVEBOOK_IDENTITY_PROVIDER` environment variable to `google_iap:<your-jwt-audience>`
 
   For more information about Google IAP,
   see: https://cloud.google.com/iap/docs/concepts-overview

test/livebook/app_test.exs

@@ -182,7 +182,7 @@ defmodule Livebook.AppTest do
       assert %{sessions: [%{id: ^session_id2}, %{id: ^session_id1}]} = App.get_by_pid(app_pid)
     end
 
-    test "adds optional user information in multi-session mode" do
+    test "does not store session creator user with personal hub" do
       slug = Utils.random_short_id()
       app_settings = %{Notebook.AppSettings.new() | slug: slug, multi_session: true}
       notebook = %{Notebook.new() | app_settings: app_settings}
@@ -194,16 +194,22 @@ defmodule Livebook.AppTest do
       user = %{Livebook.Users.User.new() | name: "Jake Peralta"}
       session_id = App.get_session_id(app_pid, user: user)
 
+      assert %{sessions: [%{id: ^session_id, started_by: nil}]} = App.get_by_pid(app_pid)
+    end
+
+    test "stores session creator user in multi-session mode with teams hub" do
+      slug = Utils.random_short_id()
+      app_settings = %{Notebook.AppSettings.new() | slug: slug, multi_session: true}
+      notebook = %{Notebook.new() | app_settings: app_settings, teams_enabled: true}
+
+      app_pid = start_app(notebook)
+
+      assert %{sessions: []} = App.get_by_pid(app_pid)
+
+      user = %{Livebook.Users.User.new() | name: "Jake Peralta"}
+      session_id = App.get_session_id(app_pid, user: user)
+
       assert %{sessions: [%{id: ^session_id, started_by: ^user}]} = App.get_by_pid(app_pid)
-
-      # Tracks user updates
-
-      App.subscribe(slug)
-
-      user = %{user | name: "Jake"}
-      Livebook.Users.update_user(user)
-
-      assert_receive {:app_updated, %{sessions: [%{id: ^session_id, started_by: ^user}]}}
     end
   end
 

test/livebook/live_markdown/import_test.exs

@@ -1149,7 +1149,11 @@ defmodule Livebook.LiveMarkdown.ImportTest do
 
       {notebook, []} = Import.notebook_from_livemd(markdown)
 
-      assert %Notebook{hub_id: "personal-hub", hub_secret_names: ["DB_PASSWORD"]} = notebook
+      assert %Notebook{
+               hub_id: "personal-hub",
+               hub_secret_names: ["DB_PASSWORD"],
+               teams_enabled: true
+             } = notebook
     end
 
     test "returns a warning when notebook stamp is invalid using offline hub" do
@@ -1178,8 +1182,33 @@ defmodule Livebook.LiveMarkdown.ImportTest do
       <!-- livebook:{"offset":58,"stamp":{"token":"invalid","token_signature":"invalid","version":1}} -->
       """
 
-      assert {%Notebook{hub_secret_names: []}, ["failed to verify notebook stamp"]} =
-               Import.notebook_from_livemd(markdown)
+      assert {%Notebook{
+                hub_id: "personal-hub",
+                hub_secret_names: [],
+                teams_enabled: false
+              }, ["failed to verify notebook stamp"]} = Import.notebook_from_livemd(markdown)
+    end
+
+    test "sets :teams_enabled to true when the teams hub exist regardless the stamp" do
+      %{id: hub_id} = Livebook.Factory.insert_hub(:team)
+
+      markdown = """
+      <!-- livebook:{"hub_id":"#{hub_id}"} -->
+
+      # My Notebook
+
+      ## Section 1
+
+      ```elixir
+      IO.puts("hey")
+      ```
+
+      <!-- livebook:{"offset":58,"stamp":{"token":"invalid","token_signature":"invalid","version":1}} -->
+      """
+
+      {notebook, [_]} = Import.notebook_from_livemd(markdown)
+
+      assert %Notebook{hub_id: ^hub_id, teams_enabled: true} = notebook
     end
   end
 end

test/livebook/session_test.exs

@@ -1301,6 +1301,94 @@ defmodule Livebook.SessionTest do
 
       App.close(app.pid)
     end
+
+    test "app session responds to app info request" do
+      slug = Utils.random_short_id()
+      app_settings = %{Notebook.AppSettings.new() | slug: slug, multi_session: true}
+      notebook = %{Notebook.new() | app_settings: app_settings}
+
+      user = Livebook.Users.User.new()
+
+      # Multi-session
+
+      {:ok, app_pid} = Apps.deploy(notebook)
+      session_id = App.get_session_id(app_pid, user: user)
+      {:ok, session} = Livebook.Sessions.fetch_session(session_id)
+
+      send(session.pid, {:runtime_app_info_request, self()})
+      assert_receive {:runtime_app_info_reply, app_info}
+
+      assert app_info == %{type: :multi_session}
+
+      # Single-session
+
+      notebook = put_in(notebook.app_settings.multi_session, false)
+      {:ok, app_pid} = Apps.deploy(notebook)
+      session_id = App.get_session_id(app_pid, user: user)
+      {:ok, session} = Livebook.Sessions.fetch_session(session_id)
+
+      send(session.pid, {:runtime_app_info_request, self()})
+      assert_receive {:runtime_app_info_reply, app_info}
+
+      assert app_info == %{type: :single_session}
+
+      App.close(app_pid)
+    end
+
+    test "app session responds to app info request when teams enabled" do
+      slug = Utils.random_short_id()
+      app_settings = %{Notebook.AppSettings.new() | slug: slug, multi_session: true}
+      notebook = %{Notebook.new() | app_settings: app_settings, teams_enabled: true}
+
+      user = %{
+        Livebook.Users.User.new()
+        | id: "1234",
+          name: "Jake Peralta",
+          email: "jperalta@example.com"
+      }
+
+      # Multi-session
+
+      {:ok, app_pid} = Apps.deploy(notebook)
+      session_id = App.get_session_id(app_pid, user: user)
+      {:ok, session} = Livebook.Sessions.fetch_session(session_id)
+
+      send(session.pid, {:runtime_app_info_request, self()})
+      assert_receive {:runtime_app_info_reply, app_info}
+
+      assert app_info == %{
+               type: :multi_session,
+               started_by: %{
+                 source: :session,
+                 id: "1234",
+                 name: "Jake Peralta",
+                 email: "jperalta@example.com"
+               }
+             }
+
+      # Single-session
+
+      notebook = put_in(notebook.app_settings.multi_session, false)
+      {:ok, app_pid} = Apps.deploy(notebook)
+      session_id = App.get_session_id(app_pid, user: user)
+      {:ok, session} = Livebook.Sessions.fetch_session(session_id)
+
+      send(session.pid, {:runtime_app_info_request, self()})
+      assert_receive {:runtime_app_info_reply, app_info}
+
+      assert app_info == %{type: :single_session}
+
+      App.close(app_pid)
+    end
+  end
+
+  test "responds to app session request when not app" do
+    session = start_session()
+
+    send(session.pid, {:runtime_app_info_request, self()})
+    assert_receive {:runtime_app_info_reply, app_info}
+
+    assert app_info == %{type: :none}
   end
 
   defp start_session(opts \\ []) do