Validate state in Teams identity callback to prevent login CSRF

Author: jonatanklosko

lib/livebook/zta/livebook_teams.ex

@@ -76,41 +76,22 @@ defmodule Livebook.ZTA.LivebookTeams do
     |> redirect(external: url)
   end
 
-  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, false) do
-      metadata = build_metadata(team.id, payload)
-      exp = System.os_time(:second) + 3 * 3600
-      :ets.insert(name, {access_token, {exp, metadata}})
-
-      {conn
-       |> put_session(:livebook_teams_access_token, access_token)
-       |> put_session(:livebook_teams_metadata_node, node())
-       |> redirect(to: conn.request_path)
-       |> halt(), metadata}
+  defp handle_request(name, conn, team, %{"teams_identity" => _} = params) do
+    if valid_auth_state?(conn, params) do
+      conn = delete_session(conn, :teams_auth_state)
+      handle_identity_callback(name, conn, team, params)
     else
-      _ ->
-        {conn
-         |> put_session(:teams_error, true)
-         |> redirect(to: conn.request_path)
-         |> halt(), nil}
+      restart_user_authentication(conn)
     end
   end
 
-  defp handle_request(_name, conn, _team, %{"teams_identity" => _, "failed_reason" => reason}) do
-    {conn
-     |> put_session(:teams_failed_reason, reason)
-     |> redirect(to: conn.request_path)
-     |> halt(), nil}
-  end
-
   defp handle_request(_name, conn, team, %{"teams_redirect" => _, "redirect_to" => redirect_to}) do
     case Teams.Requests.create_auth_request(team) do
       {:ok, %{"authorize_uri" => authorize_uri}} ->
         uri =
           authorize_uri
           |> URI.new!()
-          |> URI.append_query("redirect_to=#{redirect_to}")
+          |> URI.append_query(URI.encode_query(%{"redirect_to" => redirect_to}))
 
         {conn
          |> redirect(external: URI.to_string(uri))
@@ -153,6 +134,38 @@ defmodule Livebook.ZTA.LivebookTeams do
     end
   end
 
+  defp handle_identity_callback(name, conn, team, %{"code" => code}) do
+    with {:ok, access_token} <- retrieve_access_token(team, code),
+         {: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}})
+
+      {conn
+       |> put_session(:livebook_teams_access_token, access_token)
+       |> put_session(:livebook_teams_metadata_node, node())
+       |> redirect(to: conn.request_path)
+       |> halt(), metadata}
+    else
+      _ ->
+        {conn
+         |> put_session(:teams_error, true)
+         |> redirect(to: conn.request_path)
+         |> halt(), nil}
+    end
+  end
+
+  defp handle_identity_callback(_name, conn, _team, %{"failed_reason" => reason}) do
+    {conn
+     |> put_session(:teams_failed_reason, reason)
+     |> redirect(to: conn.request_path)
+     |> halt(), nil}
+  end
+
+  defp handle_identity_callback(_name, conn, _team, _params) do
+    restart_user_authentication(conn)
+  end
+
   defp retrieve_access_token(team, code) do
     with {:ok, %{"access_token" => access_token}} <-
            Teams.Requests.retrieve_access_token(team, code) do
@@ -161,6 +174,13 @@ defmodule Livebook.ZTA.LivebookTeams do
   end
 
   defp request_user_authentication(conn) do
+    # The state binds the authentication flow to this browser session,
+    # so that we only accept a code that we asked Livebook Teams for.
+    # Otherwise anyone could get a code for their own identity and have
+    # the browser complete the flow with it, effectively signing the
+    # user into someone else's account.
+    state = Livebook.Utils.random_long_id()
+
     # We have the browser do the redirect because the browser
     # knows the current page location. Unfortunately, it is quite
     # complex to know the actual host on the server, because the
@@ -174,7 +194,8 @@ defmodule Livebook.ZTA.LivebookTeams do
         <title>Redirecting...</title>
         <script>
           const redirectTo = new URL(window.location.href);
-          redirectTo.searchParams.append("teams_identity", "");
+          redirectTo.searchParams.set("teams_identity", "");
+          redirectTo.searchParams.set("teams_state", "#{state}");
 
           const url = new URL(window.location.href);
           url.searchParams.set("redirect_to", redirectTo.toString());
@@ -186,7 +207,23 @@ defmodule Livebook.ZTA.LivebookTeams do
     </html>
     """
 
-    {conn |> html(html_document) |> halt(), nil}
+    {conn |> put_session(:teams_auth_state, state) |> html(html_document) |> halt(), nil}
+  end
+
+  defp valid_auth_state?(conn, params) do
+    with state when is_binary(state) <- get_session(conn, :teams_auth_state),
+         param when is_binary(param) <- params["teams_state"] do
+      Plug.Crypto.secure_compare(state, param)
+    else
+      _ -> false
+    end
+  end
+
+  defp restart_user_authentication(conn) do
+    # We redirect instead of rendering the authentication page right
+    # away, so that the parameters of the stale callback are not
+    # carried over into the new flow
+    {conn |> redirect(to: conn.request_path) |> halt(), nil}
   end
 
   defp validate_access_token(name, conn, team, access_token) do

test/livebook_teams/zta/livebook_teams_test.exs

@@ -30,30 +30,34 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
       {conn, nil} = LivebookTeams.authenticate(test, conn, [])
       assert html_response(conn, 200) =~ "teams_redirect"
 
+      session = get_session(conn)
+      assert state = session["teams_auth_state"]
+
       redirect_to =
         LivebookWeb.Endpoint.url()
         |> URI.new!()
-        |> URI.append_query("teams_identity")
+        |> URI.append_query(URI.encode_query(%{"teams_identity" => "", "teams_state" => state}))
+        |> URI.to_string()
 
       # Step 2: Checks if the given request belongs to a browser
       conn =
-        build_conn(:get, "/", %{teams_redirect: "", redirect_to: URI.to_string(redirect_to)})
-        |> init_test_session(%{})
+        build_conn(:get, "/", %{teams_redirect: "", redirect_to: redirect_to})
+        |> init_test_session(session)
 
       {conn, nil} = LivebookTeams.authenticate(test, conn, [])
 
       # Step 3: Get redirected to Livebook Teams
       location = Phoenix.ConnTest.redirected_to(conn)
       uri = URI.parse(location)
       assert uri.path == "/identity/authorize"
-      assert %{"token" => token} = URI.decode_query(uri.query)
+      assert %{"token" => token, "redirect_to" => ^redirect_to} = URI.decode_query(uri.query)
 
       %{code: code} = TeamsRPC.allow_auth_request(node, token)
 
       # Step 4: Emulate the redirect back with the code for validation
       conn =
-        build_conn(:get, "/", %{teams_identity: "", code: code})
-        |> init_test_session(Plug.Conn.get_session(conn))
+        build_conn(:get, "/", %{teams_identity: "", teams_state: state, code: code})
+        |> init_test_session(session)
 
       assert {conn, %{id: _id, name: _, email: _, payload: %{}} = metadata} =
                LivebookTeams.authenticate(test, conn, [])
@@ -68,21 +72,82 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
       assert {%{halted: false}, ^metadata} = LivebookTeams.authenticate(test, conn, [])
     end
 
+    test "does not accept a code obtained in another authentication flow",
+         %{conn: conn, node: node, test: test} do
+      # Someone goes through the authentication flow themselves, up to
+      # the point where they have a code for their own identity
+      attacker_conn = init_test_session(conn, %{})
+      {attacker_conn, nil} = LivebookTeams.authenticate(test, attacker_conn, [])
+      assert attacker_state = get_session(attacker_conn, :teams_auth_state)
+
+      redirect_to =
+        LivebookWeb.Endpoint.url()
+        |> URI.new!()
+        |> URI.append_query(
+          URI.encode_query(%{"teams_identity" => "", "teams_state" => attacker_state})
+        )
+
+      attacker_conn =
+        build_conn(:get, "/", %{teams_redirect: "", redirect_to: URI.to_string(redirect_to)})
+        |> init_test_session(get_session(attacker_conn))
+
+      {attacker_conn, nil} = LivebookTeams.authenticate(test, attacker_conn, [])
+
+      uri = attacker_conn |> Phoenix.ConnTest.redirected_to() |> URI.parse()
+      assert %{"token" => token} = URI.decode_query(uri.query)
+
+      %{code: code} = TeamsRPC.allow_auth_request(node, token)
+
+      # Meanwhile the victim visits Livebook and starts their own flow
+      victim_conn = init_test_session(conn, %{})
+      {victim_conn, nil} = LivebookTeams.authenticate(test, victim_conn, [])
+      victim_session = get_session(victim_conn)
+
+      # Making the victim's browser complete the flow with the code has
+      # no effect, no matter which state it is presented with
+      for params <- [
+            %{teams_identity: "", teams_state: attacker_state, code: code},
+            %{teams_identity: "", code: code}
+          ] do
+        conn =
+          build_conn(:get, "/", params)
+          |> init_test_session(victim_session)
+
+        assert {conn, nil} = LivebookTeams.authenticate(test, conn, [])
+        assert redirected_to(conn, 302) == "/"
+        refute get_session(conn, :livebook_teams_access_token)
+
+        # Following the redirect starts the authentication flow over
+        conn =
+          build_conn(:get, "/")
+          |> init_test_session(get_session(conn))
+
+        assert {conn, nil} = LivebookTeams.authenticate(test, conn, [])
+        assert html_response(conn, 200) =~ "teams_redirect"
+      end
+    end
+
     test "shows an error when the user does not belong to the org", %{conn: conn, test: test} do
-      # Step 1: Emulate a request coming from Teams saying the user does belong to the org
+      # Step 1: Start the authentication flow, which stores the state in the session
       conn = init_test_session(conn, %{})
+      {conn, nil} = LivebookTeams.authenticate(test, conn, [])
+      assert state = get_session(conn, :teams_auth_state)
 
+      # Step 2: Emulate a request coming from Teams saying the user does not belong to the org
       params_from_teams = %{
         "teams_identity" => "",
+        "teams_state" => state,
         "failed_reason" => "you do not belong to this org"
       }
 
-      conn = %{conn | params: params_from_teams}
+      conn =
+        build_conn(:get, "/", params_from_teams)
+        |> init_test_session(get_session(conn))
 
       {conn, nil} = LivebookTeams.authenticate(test, conn, [])
       assert conn.status == 302
 
-      # Step 2: follow the redirect keeping the session set in previous request
+      # Step 3: follow the redirect keeping the session set in previous request
       conn =
         build_conn(:get, redirected_to(conn))
         |> init_test_session(get_session(conn))
@@ -93,6 +158,40 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
                "Failed to authenticate with Livebook Teams: you do not belong to this org"
     end
 
+    test "starts over when the callback carries neither code nor failure reason",
+         %{conn: conn, test: test} do
+      conn = init_test_session(conn, %{})
+      {conn, nil} = LivebookTeams.authenticate(test, conn, [])
+      assert state = get_session(conn, :teams_auth_state)
+
+      conn =
+        build_conn(:get, "/", %{teams_identity: "", teams_state: state})
+        |> init_test_session(get_session(conn))
+
+      {conn, nil} = LivebookTeams.authenticate(test, conn, [])
+      assert redirected_to(conn, 302) == "/"
+    end
+
+    test "ignores a failure reason from an unknown authentication flow",
+         %{conn: conn, test: test} do
+      conn = init_test_session(conn, %{})
+      {conn, nil} = LivebookTeams.authenticate(test, conn, [])
+
+      params_from_teams = %{
+        "teams_identity" => "",
+        "teams_state" => "invalid",
+        "failed_reason" => "you do not belong to this org"
+      }
+
+      conn =
+        build_conn(:get, "/", params_from_teams)
+        |> init_test_session(get_session(conn))
+
+      {conn, nil} = LivebookTeams.authenticate(test, conn, [])
+      assert redirected_to(conn, 302) == "/"
+      refute get_session(conn, :teams_failed_reason)
+    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)

test/support/integration/teams_tests.ex

@@ -186,10 +186,20 @@ defmodule Livebook.TeamsIntegrationHelper do
         "Mozilla/5.0 (X11; Linux x86_64; rv:121.0) Gecko/20100101 Firefox/121.0"
       )
 
+    # The page that starts the flow stores the state in the session,
+    # which we then need to pass back, as the browser would
+    session =
+      conn
+      |> LivebookWeb.ConnCase.with_authorization(team.id, name)
+      |> get("/")
+      |> Plug.Conn.get_session()
+
+    state = session["teams_auth_state"]
+
     redirect_to =
       LivebookWeb.Endpoint.url()
       |> URI.new!()
-      |> URI.append_query("teams_identity")
+      |> URI.append_query(URI.encode_query(%{"teams_identity" => "", "teams_state" => state}))
 
     uri =
       conn
@@ -203,8 +213,8 @@ defmodule Livebook.TeamsIntegrationHelper do
 
     session =
       conn
-      |> LivebookWeb.ConnCase.with_authorization(team.id, name)
-      |> get("/", %{teams_identity: "", code: code})
+      |> Plug.Test.init_test_session(session)
+      |> get("/", %{teams_identity: "", teams_state: state, code: code})
       |> Plug.Conn.get_session()
 
     authenticated_conn = Plug.Test.init_test_session(conn, session)