[Bugfix] Ensure the org switch flow works with redirect (#327)

* Ensure the org switch flow works with redirect

* Add the changeset

Author: ikovic

.changeset/tricky-gifts-compare.md

@@ -0,0 +1,5 @@
+---
+"@slashid/react": patch
+---
+
+Fix the issue with HRD and SSO in redirect mode

packages/react/src/context/slash-id-context.tsx

@@ -30,6 +30,14 @@ import {
   createEventBuffer,
   EventBuffer,
 } from "../components/form/event-buffer";
+import {
+  clearOrgSwitchingFlag,
+  getStorageKeyForOrgSwitchingUser,
+  LEGACY_STORAGE_TOKEN_KEY,
+  raiseOrgSwitchingFlag,
+  shouldResumeOrgSwitchingFlow,
+  STORAGE_TOKEN_KEY,
+} from "../domain/org";
 
 export type StorageOption = "memory" | "localStorage" | "cookie";
 
@@ -136,11 +144,6 @@ export const SlashIDContext =
   createContext<ISlashIDContext>(initialContextValue);
 SlashIDContext.displayName = "SlashIDContext";
 
-export const LEGACY_STORAGE_TOKEN_KEY = "@slashid/USER_TOKEN";
-
-export const STORAGE_TOKEN_KEY = (oid: string) =>
-  `${LEGACY_STORAGE_TOKEN_KEY}/${oid}`;
-
 const createStorage = (storageType: StorageOption) => {
   switch (storageType) {
     case "memory":
@@ -225,7 +228,11 @@ export const SlashIDProvider = ({
       }
 
       setUser(newUser);
-      storageRef.current?.setItem(currentOrgStorageTokenKey, newUser.token);
+
+      storageRef.current?.setItem(
+        getStorageKeyForOrgSwitchingUser(newUser) || currentOrgStorageTokenKey,
+        newUser.token
+      );
     },
     [state, currentOrgStorageTokenKey]
   );
@@ -254,6 +261,7 @@ export const SlashIDProvider = ({
       if (isNewOidTokenValid) {
         newToken = newOidToken;
       } else {
+        raiseOrgSwitchingFlag();
         newToken = await user.getTokenForOrganization(newOid);
       }
 
@@ -264,6 +272,7 @@ export const SlashIDProvider = ({
       setToken(newToken);
       setOid(newOid);
       setOrgSwitchingState({ state: "idle" });
+      clearOrgSwitchingFlag();
 
       return new User(newToken, sidRef.current);
     },
@@ -621,7 +630,11 @@ export const SlashIDProvider = ({
       ],
       {
         until: (value) => value !== null,
-        then: () => {
+        then: (foundUser) => {
+          if (foundUser && shouldResumeOrgSwitchingFlow(oid, foundUser)) {
+            __switchOrganizationInContext({ oid: foundUser.oid });
+          }
+
           setState("ready");
         },
       }
@@ -636,6 +649,8 @@ export const SlashIDProvider = ({
     storeUser,
     token,
     validateToken,
+    oid,
+    __switchOrganizationInContext,
   ]);
 
   const contextValue = useMemo<ISlashIDContext>(() => {

packages/react/src/domain/org.ts

@@ -0,0 +1,49 @@
+import { AnonymousUser, User } from "@slashid/slashid";
+
+export const LEGACY_STORAGE_TOKEN_KEY = "@slashid/USER_TOKEN";
+
+export const STORAGE_TOKEN_KEY = (oid: string) =>
+  `${LEGACY_STORAGE_TOKEN_KEY}/${oid}`;
+
+export const ORG_SWITCHING_FLAG_KEY = "@slashid/ORG_SWITCHING_FLOW";
+
+export function getStorageKeyForOrgSwitchingUser(user: User | AnonymousUser) {
+  if (sessionStorage && sessionStorage.getItem(ORG_SWITCHING_FLAG_KEY)) {
+    return STORAGE_TOKEN_KEY(user.oid);
+  }
+}
+
+/**
+ * When using SSO in org switching flows with Home Realm Discovery,
+ * redirect UX mode will cause the SDK not to switch to the proper org upon returning to the redirect URL.
+ *
+ * To fix that, we raise a flag in session storage and ensure the flow is done once we load the page again.
+ */
+export function raiseOrgSwitchingFlag() {
+  if (sessionStorage) {
+    sessionStorage.setItem(ORG_SWITCHING_FLAG_KEY, "in_progress");
+  }
+}
+
+export function clearOrgSwitchingFlag() {
+  if (sessionStorage) {
+    sessionStorage.removeItem(ORG_SWITCHING_FLAG_KEY);
+  }
+}
+
+export function shouldResumeOrgSwitchingFlow(
+  oid: string,
+  user: User | AnonymousUser | undefined | null
+) {
+  if (
+    sessionStorage &&
+    sessionStorage.getItem(ORG_SWITCHING_FLAG_KEY) &&
+    user &&
+    oid !== user.oid
+  ) {
+    return true;
+  }
+
+  clearOrgSwitchingFlag();
+  return false;
+}

packages/react/src/tests/anonymous-users.test.tsx

@@ -4,7 +4,7 @@ import { renderHook, waitFor } from "@testing-library/react";
 import { useSlashID, SlashIDProvider } from "../main";
 import { createAnonymousTestUser } from "../components/test-utils";
 import { SlashID } from "@slashid/slashid";
-import { STORAGE_TOKEN_KEY } from "../context/slash-id-context";
+import { STORAGE_TOKEN_KEY } from "../domain/org";
 
 describe("Anonymous users", () => {
   afterEach(() => {

packages/react/src/tests/token-storage.test.tsx

@@ -7,11 +7,9 @@ import {
   useSlashID,
 } from "../main";
 import { createTestUser, TEST_ORG_ID } from "../components/test-utils";
-import {
-  LEGACY_STORAGE_TOKEN_KEY,
-  STORAGE_TOKEN_KEY,
-} from "../context/slash-id-context";
+
 import userEvent from "@testing-library/user-event";
+import { STORAGE_TOKEN_KEY, LEGACY_STORAGE_TOKEN_KEY } from "../domain/org";
 
 describe("token storage key with org id suffix", () => {
   const getItemSpy = vi.spyOn(Storage.prototype, "getItem");