Fix ets deletion when token is invalid/expired (#3164)

Author: aleDsz

lib/livebook/zta/livebook_teams.ex

@@ -169,19 +169,19 @@ defmodule Livebook.ZTA.LivebookTeams do
   defp validate_access_token(name, conn, team, access_token) do
     node = get_session(conn, :livebook_teams_metadata_node)
 
+    entry =
+      try do
+        :erpc.call(node, :ets, :lookup_element, [name, access_token, 2, nil])
+      catch
+        _, _ -> nil
+      end
+
     case Teams.Requests.get_user_info(team, access_token) do
       {:ok, payload} ->
         {conn, build_metadata(team.id, payload)}
 
       :econnrefused ->
-        data =
-          try do
-            :erpc.call(node, :ets, :lookup_element, [name, access_token, 2, nil])
-          catch
-            _, _ -> nil
-          end
-
-        case {System.os_time(:second), data} do
+        case {System.os_time(:second), entry} do
           {current_timestamp, {exp, metadata}} when current_timestamp <= exp ->
             {conn, metadata}
 
@@ -196,7 +196,7 @@ defmodule Livebook.ZTA.LivebookTeams do
         end
 
       _otherwise ->
-        :ets.delete(__MODULE__, access_token)
+        entry && :erpc.call(node, :ets, :delete, [name, access_token])
         request_user_authentication(conn)
     end
   end

test/livebook_teams/zta/livebook_teams_test.exs

@@ -92,6 +92,20 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
       assert html_response(conn, 403) =~
                "Failed to authenticate with Livebook Teams: you do not belong to this org"
     end
+
+    test "deletes the cache if access token is invalid",
+         %{test: test, node: node, team: team} do
+      {conn, code} = authenticate_user_on_teams(test, node, team)
+      access_token = get_session(conn, :livebook_teams_access_token)
+      metadata_node = get_session(conn, :livebook_teams_metadata_node)
+
+      TeamsRPC.revoke_auth_request(node, code)
+      assert :erpc.call(metadata_node, :ets, :lookup_element, [test, access_token, 2, nil])
+
+      assert {%{halted: true} = conn, nil} = LivebookTeams.authenticate(test, conn, [])
+      assert html_response(conn, 200) =~ "window.location.href = "
+      refute :erpc.call(metadata_node, :ets, :lookup_element, [test, access_token, 2, nil])
+    end
   end
 
   describe "logout/2" do
@@ -126,10 +140,10 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
     @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)
+    setup :livebook_teams_auth
 
+    test "uses cached version of the identity payload",
+         %{conn: conn, test: test, node: node, code: code} do
       id = conn.assigns.current_user.id
       access_token = get_session(conn, :livebook_teams_access_token)
       groups = [%{"provider_id" => "1", "group_name" => "Foo"}]

test/support/integration/teams_rpc.ex

@@ -257,6 +257,10 @@ defmodule Livebook.TeamsRPC do
     :erpc.call(node, TeamsRPC, :allow_auth_request, [token])
   end
 
+  def revoke_auth_request(node, code) do
+    :erpc.call(node, TeamsRPC, :revoke_auth_request, [code])
+  end
+
   def toggle_teams_authentication(node, deployment_group) do
     :erpc.call(node, TeamsRPC, :toggle_teams_authentication, [deployment_group])
   end