From 0851e73937dbb2e58aba5c0a25806db838540065 Mon Sep 17 00:00:00 2001
From: Alexandre de Souza <alexandre@aledsz.com.br>
Date: Mon, 2 Dec 2024 10:25:32 -0300
Subject: [PATCH] Logout users from identity provider

---
 lib/livebook/teams/requests.ex                |  2 +-
 .../components/layout_components.ex           | 14 ++++++++
 .../controllers/logout_controller.ex          | 14 ++++++++
 lib/livebook_web/controllers/logout_html.ex   |  5 +++
 .../controllers/logout_html/logout.html.heex  | 18 ++++++++++
 lib/livebook_web/live/hooks/sidebar_hook.ex   | 35 +++++++++++++++++++
 lib/livebook_web/plugs/user_plug.ex           |  8 ++++-
 lib/livebook_web/router.ex                    |  5 +++
 8 files changed, 99 insertions(+), 2 deletions(-)
 create mode 100644 lib/livebook_web/controllers/logout_controller.ex
 create mode 100644 lib/livebook_web/controllers/logout_html.ex
 create mode 100644 lib/livebook_web/controllers/logout_html/logout.html.heex

diff --git a/lib/livebook/teams/requests.ex b/lib/livebook/teams/requests.ex
index 04631b4ca..4dfe77ac8 100644
--- a/lib/livebook/teams/requests.ex
+++ b/lib/livebook/teams/requests.ex
@@ -244,7 +244,7 @@ defmodule Livebook.Teams.Requests do
   @spec logout_identity_provider(Team.t(), String.t()) ::
           {:ok, String.t()} | {:error, map() | String.t()} | {:transport_error, String.t()}
   def logout_identity_provider(team, access_token) do
-    delete("/api/v1/org/identity/logout", %{access_token: access_token}, team)
+    post("/api/v1/org/identity/revoke", %{access_token: access_token}, team)
   end
 
   @doc """
diff --git a/lib/livebook_web/components/layout_components.ex b/lib/livebook_web/components/layout_components.ex
index d33fc1335..c403ccfc5 100644
--- a/lib/livebook_web/components/layout_components.ex
+++ b/lib/livebook_web/components/layout_components.ex
@@ -126,6 +126,20 @@ defmodule LivebookWeb.LayoutComponents do
               Shut Down
             </span>
           </button>
+          <button
+            :if={@current_user.email}
+            class="h-7 flex items-center text-gray-400 hover:text-white border-l-4 border-transparent hover:border-white"
+            aria-label="logout"
+            phx-click="logout"
+          >
+            <.remix_icon
+              icon="logout-box-line"
+              class="text-lg leading-6 w-[56px] flex justify-center"
+            />
+            <span class="text-sm font-medium">
+              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/logout_controller.ex b/lib/livebook_web/controllers/logout_controller.ex
new file mode 100644
index 000000000..ed0009cd3
--- /dev/null
+++ b/lib/livebook_web/controllers/logout_controller.ex
@@ -0,0 +1,14 @@
+defmodule LivebookWeb.LogoutController do
+  use LivebookWeb, :controller
+
+  def logout(conn, _params) do
+    if get_session(conn, :user_id) do
+      conn
+      |> configure_session(renew: true)
+      |> clear_session()
+      |> render("logout.html")
+    else
+      redirect(conn, to: ~p"/")
+    end
+  end
+end
diff --git a/lib/livebook_web/controllers/logout_html.ex b/lib/livebook_web/controllers/logout_html.ex
new file mode 100644
index 000000000..2b2508441
--- /dev/null
+++ b/lib/livebook_web/controllers/logout_html.ex
@@ -0,0 +1,5 @@
+defmodule LivebookWeb.LogoutHTML do
+  use LivebookWeb, :html
+
+  embed_templates "logout_html/*"
+end
diff --git a/lib/livebook_web/controllers/logout_html/logout.html.heex b/lib/livebook_web/controllers/logout_html/logout.html.heex
new file mode 100644
index 000000000..36994f801
--- /dev/null
+++ b/lib/livebook_web/controllers/logout_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..2d766ac54 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,23 @@ 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 ->
+        Livebook.Users.unsubscribe(socket.assigns.current_user.id)
+
+        {:halt,
+         socket
+         |> assign(current_user: nil)
+         |> redirect(to: ~p"/logout")}
+
+      :error ->
+        {:cont, socket}
+    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 +79,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 695a47102..67e20b8fe 100644
--- a/lib/livebook_web/plugs/user_plug.ex
+++ b/lib/livebook_web/plugs/user_plug.ex
@@ -94,7 +94,13 @@ 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 =
+      if identity_data = session["identity_data"] do
+        Map.new(identity_data, fn {k, v} -> {Atom.to_string(k), v} end)
+      else
+        %{}
+      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..1860e466b 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", LogoutController, :logout
+  end
+
   defp within_iframe_secure_headers(conn, _opts) do
     if Livebook.Config.within_iframe?() do
       delete_resp_header(conn, "x-frame-options")