From 601e93ae76726e4ee873da7ddb614f7c8afbe5e6 Mon Sep 17 00:00:00 2001
From: Alexandre de Souza <alexandre@aledsz.com.br>
Date: Wed, 8 Jan 2025 15:55:44 -0300
Subject: [PATCH] Implement the logout button (#2906)

---
 lib/livebook/config.ex                        | 17 ++++--
 lib/livebook/teams/requests.ex                |  9 ++++
 lib/livebook/zta/basic_auth.ex                |  2 +-
 lib/livebook/zta/livebook_teams.ex            | 11 ++++
 lib/livebook/zta/pass_through.ex              |  2 +-
 .../components/layout_components.ex           | 28 ++++++++--
 .../controllers/auth_controller.ex            | 11 ++++
 .../controllers/auth_html/logout.html.heex    | 18 +++++++
 lib/livebook_web/live/hooks/sidebar_hook.ex   | 27 ++++++++++
 lib/livebook_web/plugs/user_plug.ex           |  4 +-
 lib/livebook_web/router.ex                    |  5 ++
 .../zta/livebook_teams_test.exs               | 54 +++++++++++++++++++
 12 files changed, 177 insertions(+), 11 deletions(-)
 create mode 100644 lib/livebook_web/controllers/auth_html/logout.html.heex

diff --git a/lib/livebook/config.ex b/lib/livebook/config.ex
index a5c8e392a..367e8569e 100644
--- a/lib/livebook/config.ex
+++ b/lib/livebook/config.ex
@@ -1,8 +1,6 @@
 defmodule Livebook.Config do
   alias Livebook.FileSystem
 
-  @type authentication_mode :: :token | :password | :disabled
-
   @type authentication ::
           %{mode: :password, secret: String.t()}
           | %{mode: :token, secret: String.t()}
@@ -68,7 +66,7 @@ defmodule Livebook.Config do
   @doc """
   Returns the authentication configuration.
   """
-  @spec authentication() :: authentication_mode()
+  @spec authentication() :: authentication()
   def authentication() do
     case Application.fetch_env!(:livebook, :authentication) do
       {:password, password} -> %{mode: :password, secret: password}
@@ -270,6 +268,19 @@ defmodule Livebook.Config do
     module not in @identity_provider_no_id
   end
 
+  @doc """
+  Returns if the identity provider supports logout.
+  """
+  @spec logout_enabled?() :: boolean()
+  def logout_enabled?() do
+    {_type, module, _key} = Livebook.Config.identity_provider()
+
+    identity_logout? =
+      Code.ensure_loaded?(module) and function_exported?(module, :logout, 2)
+
+    authentication().mode != :disabled or identity_logout?
+  end
+
   @doc """
   Returns whether the application is running inside an iframe.
   """
diff --git a/lib/livebook/teams/requests.ex b/lib/livebook/teams/requests.ex
index 125a2e52b..1b3cdb435 100644
--- a/lib/livebook/teams/requests.ex
+++ b/lib/livebook/teams/requests.ex
@@ -238,6 +238,15 @@ defmodule Livebook.Teams.Requests do
     get("/api/v1/org/identity", %{access_token: access_token}, team)
   end
 
+  @doc """
+  Send a request to Livebook Team API to revoke session from given access token.
+  """
+  @spec logout_identity_provider(Team.t(), String.t()) ::
+          {:ok, String.t()} | {:error, map()} | {:transport_error, String.t()}
+  def logout_identity_provider(team, access_token) do
+    post("/api/v1/org/identity/revoke", %{access_token: access_token}, team)
+  end
+
   @doc """
   Normalizes errors map into errors for the given schema.
   """
diff --git a/lib/livebook/zta/basic_auth.ex b/lib/livebook/zta/basic_auth.ex
index 606ed2d11..a07e59f23 100644
--- a/lib/livebook/zta/basic_auth.ex
+++ b/lib/livebook/zta/basic_auth.ex
@@ -16,7 +16,7 @@ defmodule Livebook.ZTA.BasicAuth do
   end
 
   @impl true
-  def authenticate(name, conn, _options) do
+  def authenticate(name, conn, _opts) do
     {username, password} = Livebook.ZTA.get(name)
     conn = Plug.BasicAuth.basic_auth(conn, username: username, password: password)
 
diff --git a/lib/livebook/zta/livebook_teams.ex b/lib/livebook/zta/livebook_teams.ex
index 95ae2e0b3..69411b3de 100644
--- a/lib/livebook/zta/livebook_teams.ex
+++ b/lib/livebook/zta/livebook_teams.ex
@@ -35,6 +35,17 @@ defmodule Livebook.ZTA.LivebookTeams do
     end
   end
 
+  # Our extension to Livebook.ZTA to deal with logouts
+  def logout(name, %{assigns: %{current_user: %{payload: %{"access_token" => token}}}}) do
+    team = Livebook.ZTA.get(name)
+
+    case Teams.Requests.logout_identity_provider(team, token) do
+      {:ok, _no_content} -> :ok
+      {:error, %{}} -> {:error, "You are already logged out."}
+      {:transport_error, reason} -> {:error, reason}
+    end
+  end
+
   defp handle_request(conn, team, %{"teams_identity" => _, "code" => code}) do
     with {:ok, access_token} <- retrieve_access_token(team, code),
          {:ok, metadata} <- get_user_info(team, access_token) do
diff --git a/lib/livebook/zta/pass_through.ex b/lib/livebook/zta/pass_through.ex
index 77a6a9cc4..4e610410d 100644
--- a/lib/livebook/zta/pass_through.ex
+++ b/lib/livebook/zta/pass_through.ex
@@ -7,7 +7,7 @@ defmodule Livebook.ZTA.PassThrough do
   end
 
   @impl true
-  def authenticate(_, conn, _) do
+  def authenticate(_name, conn, _opts) do
     {conn, %{}}
   end
 end
diff --git a/lib/livebook_web/components/layout_components.ex b/lib/livebook_web/components/layout_components.ex
index 557ece315..e2805b9d6 100644
--- a/lib/livebook_web/components/layout_components.ex
+++ b/lib/livebook_web/components/layout_components.ex
@@ -111,21 +111,39 @@ defmodule LivebookWeb.LayoutComponents do
               to={~p"/settings"}
               current={@current_page}
             />
+            <button
+              :if={Livebook.Config.shutdown_callback()}
+              class="h-7 flex items-center text-gray-400 hover:text-white border-l-4 border-transparent hover:border-white"
+              aria-label="shutdown"
+              phx-click="shutdown"
+            >
+              <.remix_icon
+                icon="shut-down-line"
+                class="text-lg leading-6 w-[56px] flex justify-center"
+              />
+              <span class="text-sm font-medium">
+                Shut Down
+              </span>
+            </button>
           </div>
           <.hub_section hubs={@saved_hubs} current_page={@current_page} />
         </div>
         <div class="flex flex-col">
           <button
-            :if={Livebook.Config.shutdown_callback()}
+            :if={Livebook.Config.logout_enabled?()}
             class="h-7 flex items-center text-gray-400 hover:text-white border-l-4 border-transparent hover:border-white"
-            aria-label="shutdown"
-            phx-click="shutdown"
+            aria-label="logout"
+            phx-click="logout"
           >
-            <.remix_icon icon="shut-down-line" class="text-lg leading-6 w-[56px] flex justify-center" />
+            <.remix_icon
+              icon="logout-box-line"
+              class="text-lg leading-6 w-[56px] flex justify-center"
+            />
             <span class="text-sm font-medium">
-              Shut Down
+              Logout
             </span>
           </button>
+
           <button
             class="mt-6 flex items-center group border-l-4 border-transparent"
             aria_label="user profile"
diff --git a/lib/livebook_web/controllers/auth_controller.ex b/lib/livebook_web/controllers/auth_controller.ex
index c62b06805..f067a4bb5 100644
--- a/lib/livebook_web/controllers/auth_controller.ex
+++ b/lib/livebook_web/controllers/auth_controller.ex
@@ -46,6 +46,17 @@ defmodule LivebookWeb.AuthController do
     end
   end
 
+  def logout(conn, _params) do
+    if get_session(conn, :user_id) do
+      conn
+      |> configure_session(renew: true)
+      |> clear_session()
+      |> render("logout.html")
+    else
+      redirect_to(conn)
+    end
+  end
+
   defp render_form_error(conn, authentication_mode) do
     errors = [{"%{authentication_mode} is invalid", [authentication_mode: authentication_mode]}]
 
diff --git a/lib/livebook_web/controllers/auth_html/logout.html.heex b/lib/livebook_web/controllers/auth_html/logout.html.heex
new file mode 100644
index 000000000..36994f801
--- /dev/null
+++ b/lib/livebook_web/controllers/auth_html/logout.html.heex
@@ -0,0 +1,18 @@
+<div class="h-screen w-full px-4 py-8 bg-gray-900 flex justify-center items-center">
+  <div class="max-w-[400px] w-full flex flex-col">
+    <a href={~p"/"} class="mb-2 -ml-2">
+      <img src={~p"/images/logo.png"} height="96" width="96" alt="livebook" />
+    </a>
+    <div class="mb-2 text-xl text-gray-100 font-medium">
+      You have been logged out
+    </div>
+
+    <div class="mb-8 text-sm text-gray-200">
+      Thank you for using <strong>Livebook</strong>
+    </div>
+
+    <div class="text-gray-50 w-full">
+      <.button navigate={~p"/"}>Sign in back</.button>
+    </div>
+  </div>
+</div>
diff --git a/lib/livebook_web/live/hooks/sidebar_hook.ex b/lib/livebook_web/live/hooks/sidebar_hook.ex
index 14f7f4b58..34f337efc 100644
--- a/lib/livebook_web/live/hooks/sidebar_hook.ex
+++ b/lib/livebook_web/live/hooks/sidebar_hook.ex
@@ -1,4 +1,5 @@
 defmodule LivebookWeb.SidebarHook do
+  use LivebookWeb, :verified_routes
   require Logger
 
   import Phoenix.Component
@@ -17,6 +18,8 @@ defmodule LivebookWeb.SidebarHook do
       |> attach_hook(:hubs, :handle_info, &handle_info/2)
       |> attach_hook(:shutdown, :handle_info, &handle_info/2)
       |> attach_hook(:shutdown, :handle_event, &handle_event/3)
+      |> attach_hook(:logout, :handle_info, &handle_info/2)
+      |> attach_hook(:logout, :handle_event, &handle_event/3)
 
     {:cont, socket}
   end
@@ -25,6 +28,15 @@ defmodule LivebookWeb.SidebarHook do
     {:halt, put_flash(socket, :info, "Livebook is shutting down. You can close this page.")}
   end
 
+  defp handle_info(:logout, socket) do
+    {_type, module, _key} = Livebook.Config.identity_provider()
+
+    case module.logout(LivebookWeb.ZTA, socket) do
+      :ok -> {:halt, redirect(socket, to: ~p"/logout")}
+      {:error, reason} -> {:cont, put_flash(socket, :error, reason)}
+    end
+  end
+
   @connection_events ~w(hub_connected hub_changed hub_deleted)a
 
   defp handle_info(event, socket) when elem(event, 0) in @connection_events do
@@ -59,5 +71,20 @@ defmodule LivebookWeb.SidebarHook do
      )}
   end
 
+  defp handle_event("logout", _params, socket) do
+    on_confirm = fn socket ->
+      Phoenix.PubSub.broadcast(Livebook.PubSub, "sidebar", :logout)
+      put_flash(socket, :info, "Livebook is logging out. You will be redirected soon.")
+    end
+
+    {:halt,
+     confirm(socket, on_confirm,
+       title: "Log out",
+       description: "Are you sure you want to log out Livebook now?",
+       confirm_text: "Log out",
+       confirm_icon: "logout-box-line"
+     )}
+  end
+
   defp handle_event(_event, _params, socket), do: {:cont, socket}
 end
diff --git a/lib/livebook_web/plugs/user_plug.ex b/lib/livebook_web/plugs/user_plug.ex
index 66a66d637..0696e59b3 100644
--- a/lib/livebook_web/plugs/user_plug.ex
+++ b/lib/livebook_web/plugs/user_plug.ex
@@ -94,7 +94,9 @@ defmodule LivebookWeb.UserPlug do
   we get possibly updated `user_data` from `connect_params`.
   """
   def build_current_user(session, user_data_override \\ nil) do
-    identity_data = Map.new(session["identity_data"], fn {k, v} -> {Atom.to_string(k), v} end)
+    identity_data =
+      Map.new(session["identity_data"] || %{}, fn {k, v} -> {Atom.to_string(k), v} end)
+
     attrs = user_data_override || session["user_data"] || %{}
 
     attrs =
diff --git a/lib/livebook_web/router.ex b/lib/livebook_web/router.ex
index 2ef3d849f..c15c3a7f4 100644
--- a/lib/livebook_web/router.ex
+++ b/lib/livebook_web/router.ex
@@ -169,6 +169,11 @@ defmodule LivebookWeb.Router do
     post "/", AuthController, :authenticate
   end
 
+  scope "/", LivebookWeb do
+    pipe_through [:browser]
+    get "/logout", AuthController, :logout
+  end
+
   defp within_iframe_secure_headers(conn, _opts) do
     if Livebook.Config.within_iframe?() do
       delete_resp_header(conn, "x-frame-options")
diff --git a/test/livebook_teams/zta/livebook_teams_test.exs b/test/livebook_teams/zta/livebook_teams_test.exs
index 4bdc97ceb..e0da77f24 100644
--- a/test/livebook_teams/zta/livebook_teams_test.exs
+++ b/test/livebook_teams/zta/livebook_teams_test.exs
@@ -100,4 +100,58 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
                "Failed to authenticate with Livebook Teams: you do not belong to this org"
     end
   end
+
+  describe "logout/2" do
+    test "revoke access token from Livebook Teams", %{conn: conn, node: node, test: test} do
+      # Step 1: Get redirected to Livebook Teams
+      conn = init_test_session(conn, %{})
+      {conn, nil} = LivebookTeams.authenticate(test, conn, [])
+
+      [_, location] = Regex.run(~r/URL\("(.*?)"\)/, html_response(conn, 200))
+      uri = URI.parse(location)
+      assert uri.path == "/identity/authorize"
+      assert %{"code" => code} = URI.decode_query(uri.query)
+
+      erpc_call(node, :allow_auth_request, [code])
+
+      # Step 2: Emulate the redirect back with the code for validation
+      conn =
+        build_conn(:get, "/", %{teams_identity: "", code: code})
+        |> init_test_session(%{})
+
+      assert {conn, %{id: _id, name: _, email: _, payload: %{"access_token" => _}} = metadata} =
+               LivebookTeams.authenticate(test, conn, [])
+
+      assert redirected_to(conn, 302) == "/"
+
+      # Step 3: Confirm the token/metadata is valid for future requests
+      conn =
+        build_conn(:get, "/")
+        |> init_test_session(%{identity_data: metadata})
+
+      assert {%{halted: false}, ^metadata} = LivebookTeams.authenticate(test, conn, [])
+
+      # Step 4: Revoke the token and the metadata will be invalid for future requests
+      user =
+        metadata.id
+        |> Livebook.Users.User.new()
+        |> Livebook.Users.User.changeset(metadata)
+        |> Ecto.Changeset.apply_changes()
+
+      conn =
+        build_conn(:get, "/")
+        |> init_test_session(%{identity_data: metadata})
+        |> assign(:current_user, user)
+
+      assert LivebookTeams.logout(test, conn) == :ok
+
+      # Step 5: If we try to revoke again, it should fail
+      assert {:error, _} = LivebookTeams.logout(test, conn)
+
+      # Step 6: It we try to authenticate again, it should redirect to Teams
+      {conn, nil} = LivebookTeams.authenticate(test, conn, [])
+      assert conn.halted
+      assert html_response(conn, 200) =~ "window.location.href = "
+    end
+  end
 end