diff --git a/lib/livebook/teams/requests.ex b/lib/livebook/teams/requests.ex index 8420ae66828..25166935512 100644 --- a/lib/livebook/teams/requests.ex +++ b/lib/livebook/teams/requests.ex @@ -225,17 +225,21 @@ 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() | :econnrefused - def get_user_info(team, access_token) do + @spec get_user_info(Team.t(), String.t(), boolean()) :: api_result() | :econnrefused + def get_user_info(team, access_token, valid_cache?) do 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 + fun = fn + _, %Req.TransportError{reason: :econnrefused} -> not valid_cache? + _, %Req.Response{status: status} when status in [408, 429, 500, 502, 503, 504] -> true + _, %Req.TransportError{reason: reason} when reason in [:timeout, :closed] -> true + _, _ -> false + end - otherwise -> - handle_response(otherwise) + case Req.get(req, url: "/api/v1/org/identity", params: params, retry: fun) do + {:error, %Req.TransportError{reason: :econnrefused}} -> :econnrefused + otherwise -> handle_response(otherwise) end end @@ -330,17 +334,6 @@ 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 diff --git a/lib/livebook/zta/livebook_teams.ex b/lib/livebook/zta/livebook_teams.ex index b35a5e38bf0..27a3674badf 100644 --- a/lib/livebook/zta/livebook_teams.ex +++ b/lib/livebook/zta/livebook_teams.ex @@ -69,7 +69,7 @@ defmodule Livebook.ZTA.LivebookTeams do defp handle_request(name, conn, team, %{"teams_identity" => _, "code" => code}) do with {:ok, access_token} <- retrieve_access_token(team, code), - {:ok, payload} <- Teams.Requests.get_user_info(team, access_token) do + {:ok, payload} <- Teams.Requests.get_user_info(team, access_token, false) do metadata = build_metadata(team.id, payload) exp = System.os_time(:second) + 3 * 3600 :ets.insert(name, {access_token, {exp, metadata}}) @@ -183,30 +183,33 @@ defmodule Livebook.ZTA.LivebookTeams do defp validate_access_token(name, conn, team, access_token) do node = get_session(conn, :livebook_teams_metadata_node) - case Teams.Requests.get_user_info(team, access_token) do + entry = + try do + :erpc.call(node, :ets, :lookup_element, [name, access_token, 2, nil]) + catch + _, _ -> nil + end + + valid_cache? = valid_cache?(entry) + + case Teams.Requests.get_user_info(team, access_token, valid_cache?) do {:ok, payload} -> {conn, build_metadata(team.id, payload)} :econnrefused -> - entry = - try do - :erpc.call(node, :ets, :lookup_element, [name, access_token, 2, nil]) - catch - _, _ -> nil - end - - case {System.os_time(:second), entry} do - {current_timestamp, {exp, metadata}} when current_timestamp <= exp -> - {conn, metadata} - - {_, entry} -> - entry && :erpc.call(node, :ets, :delete, [name, access_token]) - - {conn - |> put_status(:service_unavailable) - |> put_view(LivebookWeb.ErrorHTML) - |> render("503.html") - |> halt(), nil} + # We double-checked because the response may contain latency, + # so the timestamp must be revalidated to ensure it hasn't expired. + if valid_cache?(entry) do + {_, metadata} = entry + {conn, metadata} + else + entry && :erpc.call(node, :ets, :delete, [name, access_token]) + + {conn + |> put_status(:service_unavailable) + |> put_view(LivebookWeb.ErrorHTML) + |> render("503.html") + |> halt(), nil} end _otherwise -> @@ -220,6 +223,9 @@ defmodule Livebook.ZTA.LivebookTeams do end end + defp valid_cache?({exp, _}), do: System.os_time(:second) <= exp + defp valid_cache?(_entry), do: false + @doc """ Returns the user metadata from given payload. """ diff --git a/test/livebook_teams/zta/livebook_teams_test.exs b/test/livebook_teams/zta/livebook_teams_test.exs index 62f7a8985ae..2c7b0b8dcca 100644 --- a/test/livebook_teams/zta/livebook_teams_test.exs +++ b/test/livebook_teams/zta/livebook_teams_test.exs @@ -134,6 +134,7 @@ defmodule Livebook.ZTA.LivebookTeamsTest do alias Livebook.ZTA.LivebookTeams + @moduletag :capture_log @moduletag teams_for: :agent setup :teams @@ -142,40 +143,37 @@ defmodule Livebook.ZTA.LivebookTeamsTest do setup :livebook_teams_auth - test "uses cached version of the identity payload", - %{conn: conn, test: test, node: node, code: code} do + test "uses cached version of the identity payload", %{conn: conn, test: test} = ctx do + Application.put_env(:livebook, :teams_url, "http://localhost:1234") + 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") + node = get_session(conn, :livebook_teams_metadata_node) # update the groups, but doesn't return because Livebook is using the cached one - TeamsRPC.update_user_info_groups(node, code, groups) + TeamsRPC.update_user_info_groups(ctx.node, ctx.code, groups) + + # shouldn't retry the request + current_timestamp = System.os_time(:second) assert {_, %{id: ^id, groups: []}} = LivebookTeams.authenticate(test, conn, []) + assert System.os_time(:second) - current_timestamp < :timer.seconds(1) # simulate if the token already expired exp = System.os_time(:second) - 5 * 60 - metadata_node = get_session(conn, :livebook_teams_metadata_node) - - {_, metadata} = - :erpc.call(metadata_node, :ets, :lookup_element, [test, access_token, 2, nil]) - - :erpc.call(metadata_node, :ets, :insert, [test, {access_token, {exp, metadata}}]) - - # now it should return status 503 - assert {%{status: 503, halted: true, resp_body: body}, nil} = - LivebookTeams.authenticate(test, conn, []) + {_, metadata} = :erpc.call(node, :ets, :lookup_element, [test, access_token, 2, nil]) + :erpc.call(node, :ets, :insert, [test, {access_token, {exp, metadata}}]) - assert body =~ "The server is currently down or under maintenance" + # now it should retry to request to Teams and return status 503 + assert ExUnit.CaptureLog.capture_log(fn -> + assert {%{status: 503, halted: true, resp_body: body}, nil} = + LivebookTeams.authenticate(test, conn, []) - # still show 503 error page because Teams isn't up yet - assert {%{status: 503, halted: true}, nil} = LivebookTeams.authenticate(test, conn, []) + assert body =~ "The server is currently down or under maintenance" + end) =~ "retry: got exception, will retry in" # now gets the updated userinfo from Teams - Application.put_env(:livebook, :teams_url, url) + Application.put_env(:livebook, :teams_url, TeamsServer.url()) assert {_conn, %{id: ^id, groups: ^groups}} = LivebookTeams.authenticate(test, conn, []) end end