From 96a995bf26d852f1ac99fe136d84330a5575762f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Hugo=20Bara=C3=BAna?= <hugo.barauna@gmail.com>
Date: Wed, 20 Nov 2024 23:25:22 -0300
Subject: [PATCH] Fix auth with Livebook teams when user not belongs to org
 (#2865)

---
 lib/livebook/zta/livebook_teams.ex            |  2 +-
 .../zta/livebook_teams_test.exs               | 26 +++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/lib/livebook/zta/livebook_teams.ex b/lib/livebook/zta/livebook_teams.ex
index 80d3d3899..1adc692db 100644
--- a/lib/livebook/zta/livebook_teams.ex
+++ b/lib/livebook/zta/livebook_teams.ex
@@ -53,8 +53,8 @@ defmodule Livebook.ZTA.LivebookTeams do
 
   defp handle_request(conn, _team, %{"teams_identity" => _, "failed_reason" => reason}) do
     {conn
-     |> redirect(to: conn.request_path)
      |> put_session(:teams_failed_reason, reason)
+     |> redirect(to: conn.request_path)
      |> halt(), nil}
   end
 
diff --git a/test/livebook_teams/zta/livebook_teams_test.exs b/test/livebook_teams/zta/livebook_teams_test.exs
index 866161b90..129dab218 100644
--- a/test/livebook_teams/zta/livebook_teams_test.exs
+++ b/test/livebook_teams/zta/livebook_teams_test.exs
@@ -75,5 +75,31 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
       assert conn.halted
       assert html_response(conn, 200) =~ "window.location.href = "
     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
+      conn = init_test_session(conn, %{})
+
+      params_from_teams = %{
+        "teams_identity" => "",
+        "failed_reason" => "you do not belong to this org"
+      }
+
+      conn = %Plug.Conn{conn | params: params_from_teams}
+
+      {conn, nil} = LivebookTeams.authenticate(test, conn, [])
+      session = Plug.Conn.get_session(conn)
+
+      assert conn.status == 302
+
+      # Step 2: follow the redirect keeping the session set in previous request
+      conn = build_conn(:get, redirected_to(conn))
+      conn = init_test_session(conn, session)
+
+      {conn, nil} = LivebookTeams.authenticate(test, conn, [])
+
+      assert html_response(conn, 200) =~
+               "Failed to authenticate with Livebook Teams: you do not belong to this org"
+    end
   end
 end