Keeps the user data while Livebook Teams is down (#3156)

Author: aleDsz

lib/livebook/teams/requests.ex

@@ -225,9 +225,18 @@ defmodule Livebook.Teams.Requests do
   @doc """
   Send a request to Livebook Team API to get the user information from given access token.
   """
-  @spec get_user_info(Team.t(), String.t()) :: api_result()
+  @spec get_user_info(Team.t(), String.t()) :: api_result() | :econnrefused
   def get_user_info(team, access_token) do
-    get("/api/v1/org/identity", %{access_token: access_token}, team)
+    req = build_req(team)
+    params = %{access_token: access_token}
+
+    case Req.get(req, url: "/api/v1/org/identity", params: params) do
+      {:error, %Req.TransportError{reason: :econnrefused}} ->
+        :econnrefused
+
+      otherwise ->
+        handle_response(otherwise)
+    end
   end
 
   @doc """
@@ -321,6 +330,17 @@ defmodule Livebook.Teams.Requests do
     |> Req.Request.put_new_header("x-lb-version", Livebook.Config.app_version())
     |> Livebook.Utils.req_attach_defaults()
     |> add_team_auth(team)
+    |> put_test_req_options()
+  end
+
+  if Mix.env() == :test do
+    defp put_test_req_options(req) do
+      Req.Request.merge_options(req, retry: false)
+    end
+  else
+    defp put_test_req_options(req) do
+      req
+    end
   end
 
   defp add_team_auth(req, nil), do: req

lib/livebook/zta/livebook_teams.ex

@@ -8,21 +8,27 @@ defmodule Livebook.ZTA.LivebookTeams do
 
   @behaviour NimbleZTA
 
-  @impl true
+  @exp_timestamp_sec System.os_time(:second) + 3 * 3600
+
+  @impl NimbleZTA
   def child_spec(opts) do
     %{id: __MODULE__, start: {__MODULE__, :start_link, [opts]}}
   end
 
   def start_link(opts) do
     name = Keyword.fetch!(opts, :name)
-    identity_key = Keyword.fetch!(opts, :identity_key)
-    team = Livebook.Hubs.fetch_hub!(identity_key)
+    id = Keyword.fetch!(opts, :identity_key)
+    team = Livebook.Hubs.fetch_hub!(id)
+
+    if :ets.whereis(__MODULE__) == :undefined do
+      :ets.new(__MODULE__, [:named_table, :public, :set, read_concurrency: true])
+    end
 
     NimbleZTA.put(name, team)
     :ignore
   end
 
-  @impl true
+  @impl NimbleZTA
   def authenticate(name, conn, _opts) do
     team = NimbleZTA.get(name)
 
@@ -53,7 +59,10 @@ defmodule Livebook.ZTA.LivebookTeams do
 
   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
+         {:ok, payload} <- Teams.Requests.get_user_info(team, access_token) do
+      metadata = build_metadata(team.id, payload)
+      :ets.insert(__MODULE__, {access_token, {@exp_timestamp_sec, metadata}})
+
       {conn
        |> put_session(:livebook_teams_access_token, access_token)
        |> redirect(to: conn.request_path)
@@ -99,7 +108,6 @@ defmodule Livebook.ZTA.LivebookTeams do
       %{"livebook_teams_access_token" => access_token} ->
         validate_access_token(conn, team, access_token)
 
-      # it means, we couldn't reach to Teams server
       %{"teams_error" => true} ->
         {conn
          |> put_status(:bad_request)
@@ -124,13 +132,6 @@ defmodule Livebook.ZTA.LivebookTeams do
     end
   end
 
-  defp validate_access_token(conn, team, access_token) do
-    case get_user_info(team, access_token) do
-      {:ok, metadata} -> {conn, metadata}
-      _ -> request_user_authentication(conn)
-    end
-  end
-
   defp retrieve_access_token(team, code) do
     with {:ok, %{"access_token" => access_token}} <-
            Teams.Requests.retrieve_access_token(team, code) do
@@ -167,9 +168,31 @@ defmodule Livebook.ZTA.LivebookTeams do
     {conn |> html(html_document) |> halt(), nil}
   end
 
-  defp get_user_info(team, access_token) do
-    with {:ok, payload} <- Teams.Requests.get_user_info(team, access_token) do
-      {:ok, build_metadata(team.id, payload)}
+  defp validate_access_token(conn, team, access_token) do
+    case Teams.Requests.get_user_info(team, access_token) do
+      {:ok, payload} ->
+        {conn, build_metadata(team.id, payload)}
+
+      :econnrefused ->
+        data = :ets.lookup_element(__MODULE__, access_token, 2, nil)
+
+        case {System.os_time(:second), data} do
+          {current_timestamp, {exp, metadata}} when current_timestamp <= exp ->
+            {conn, metadata}
+
+          {_, entry} ->
+            entry && :ets.delete(__MODULE__, access_token)
+
+            {conn
+             |> put_status(:service_unavailable)
+             |> put_view(LivebookWeb.ErrorHTML)
+             |> render("503.html")
+             |> halt(), nil}
+        end
+
+      _otherwise ->
+        :ets.delete(__MODULE__, access_token)
+        request_user_authentication(conn)
     end
   end
 

lib/livebook_web/controllers/error_html.ex

@@ -25,6 +25,16 @@ defmodule LivebookWeb.ErrorHTML do
     """
   end
 
+  def render("503.html", assigns) do
+    ~H"""
+    <.error_page
+      status={503}
+      title="Service unavailable"
+      details="The server is currently down or under maintenance"
+    />
+    """
+  end
+
   def render(_template, assigns) do
     ~H"""
     <.error_page
@@ -50,7 +60,13 @@ defmodule LivebookWeb.ErrorHTML do
         <link rel="icon" type="image/svg+xml" href={~p"/favicons/favicon.svg"} />
         <link rel="alternate icon" type="image/png" href={~p"/favicons/favicon.png"} />
         <title>{@status} - Livebook</title>
-        <link rel="stylesheet" href={~p"/assets/app.css"} />
+        <%= if LivebookWeb.Layouts.dev?() do %>
+          <script phx-track-static type="module" src="http://localhost:4432/@vite/client">
+          </script>
+          <link phx-track-static rel="stylesheet" href="http://localhost:4432/css/app.css" />
+        <% else %>
+          <link phx-track-static rel="stylesheet" href={~p"/assets/app.css"} />
+        <% end %>
       </head>
       <body>
         <div class="h-screen flex items-center justify-center bg-gray-900">

test/livebook_teams/zta/livebook_teams_test.exs

@@ -113,4 +113,52 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
       assert html_response(conn, 200) =~ "window.location.href = "
     end
   end
+
+  defmodule Global do
+    # We need to "turn off" the Teams API during test
+    use Livebook.TeamsIntegrationCase, async: false
+
+    alias Livebook.ZTA.LivebookTeams
+
+    @moduletag teams_for: :agent
+    setup :teams
+
+    @moduletag subscribe_to_hubs_topics: [:connection]
+    @moduletag subscribe_to_teams_topics: [:clients, :agents]
+
+    test "uses cached version of the identity payload", %{test: test, team: team, node: node} do
+      start_supervised!({LivebookTeams, name: test, identity_key: team.id})
+      {conn, code} = authenticate_user_on_teams(test, node, team)
+
+      id = conn.assigns.current_user.id
+      access_token = get_session(conn, :livebook_teams_access_token)
+      groups = [%{"provider_id" => "1", "group_name" => "Foo"}]
+
+      # simulate the Teams API is down
+      url = Livebook.Config.teams_url()
+      Application.put_env(:livebook, :teams_url, "http://localhost:1234")
+
+      # update the groups, but doesn't return because Livebook is using the cached one
+      TeamsRPC.update_user_info_groups(node, code, groups)
+      assert {_, %{id: ^id, groups: []}} = LivebookTeams.authenticate(test, conn, [])
+
+      # simulate if the token already expired
+      exp = System.os_time(:second) - 5 * 60
+      {_, metadata} = :ets.lookup_element(LivebookTeams, access_token, 2)
+      :ets.insert(LivebookTeams, {access_token, {exp, metadata}})
+
+      # now it should return status 503
+      assert {%{status: 503, halted: true, resp_body: body}, nil} =
+               LivebookTeams.authenticate(test, conn, [])
+
+      assert body =~ "The server is currently down or under maintenance"
+
+      # still show 503 error page because Teams isn't up yet
+      assert {%{status: 503, halted: true}, nil} = LivebookTeams.authenticate(test, conn, [])
+
+      # now gets the updated userinfo from Teams
+      Application.put_env(:livebook, :teams_url, url)
+      assert {_conn, %{id: ^id, groups: ^groups}} = LivebookTeams.authenticate(test, conn, [])
+    end
+  end
 end

test/support/teams_integration_case.ex

@@ -27,7 +27,7 @@ defmodule Livebook.TeamsIntegrationCase do
     url = TeamsServer.url()
     node = TeamsServer.get_node()
 
-    Application.put_env(:livebook, :teams_url, url, persistent: true)
+    Application.put_env(:livebook, :teams_url, url)
 
     {:ok, node: node}
   end