Implement the logout button

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 get the user information from given access token.
+  """
+  @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
+    post("/api/v1/org/identity/revoke", %{access_token: access_token}, team)
+  end
+
   @doc """
   Normalizes errors map into errors for the given schema.
   """

lib/livebook/zta.ex

@@ -45,6 +45,11 @@ defmodule Livebook.ZTA do
   """
   @callback authenticate(name(), Plug.Conn.t(), keyword()) :: {Plug.Conn.t(), metadata() | nil}
 
+  @doc """
+  Logouts against the given name.
+  """
+  @callback logout(name(), Phoenix.LiveView.Socket.t()) :: :ok | :error
+
   @doc false
   def init do
     :ets.new(__MODULE__, [:named_table, :public, :set, read_concurrency: true])

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)
 
@@ -26,4 +26,9 @@ defmodule Livebook.ZTA.BasicAuth do
       {conn, %{}}
     end
   end
+
+  @impl true
+  def logout(_name, _socket) do
+    :error
+  end
 end

lib/livebook/zta/cloudflare.ex

@@ -25,6 +25,11 @@ defmodule Livebook.ZTA.Cloudflare do
     {conn, authenticate_user(token, identity, keys)}
   end
 
+  @impl true
+  def logout(_name, _socket) do
+    :error
+  end
+
   @impl true
   def init(options) do
     state = struct!(__MODULE__, options)

lib/livebook/zta/google_iap.ex

@@ -25,6 +25,11 @@ defmodule Livebook.ZTA.GoogleIAP do
     {conn, authenticate_user(token, identity, keys)}
   end
 
+  @impl true
+  def logout(_name, _socket) do
+    :error
+  end
+
   @impl true
   def init(options) do
     state = struct!(__MODULE__, options)

lib/livebook/zta/livebook_teams.ex

@@ -35,6 +35,16 @@ defmodule Livebook.ZTA.LivebookTeams do
     end
   end
 
+  @impl true
+  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
+      _otherwise -> :error
+    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

lib/livebook/zta/pass_through.ex

@@ -7,7 +7,12 @@ defmodule Livebook.ZTA.PassThrough do
   end
 
   @impl true
-  def authenticate(_, conn, _) do
+  def authenticate(_name, conn, _opts) do
     {conn, %{}}
   end
+
+  @impl true
+  def logout(_name, _socket) do
+    :error
+  end
 end

lib/livebook/zta/tailscale.ex

@@ -29,6 +29,11 @@ defmodule Livebook.ZTA.Tailscale do
     {conn, user}
   end
 
+  @impl true
+  def logout(_name, _socket) do
+    :error
+  end
+
   defp authenticate_ip(remote_ip, address) do
     {url, options} =
       if String.starts_with?(address, "http") do

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"

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

lib/livebook_web/controllers/logout_html.ex

@@ -0,0 +1,5 @@
+defmodule LivebookWeb.LogoutHTML do
+  use LivebookWeb, :html
+
+  embed_templates "logout_html/*"
+end

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>

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

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 =

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")

test/livebook_teams/zta/livebook_teams_test.exs

@@ -100,4 +100,50 @@ 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
+    end
+  end
 end