feat(relocation): Add auto redirect after 2FA enroll (#68645)

Author: azaslavsky

static/app/views/settings/account/accountSecurity/accountSecurityEnroll.spec.tsx

@@ -1,14 +1,25 @@
 import {AuthenticatorsFixture} from 'sentry-fixture/authenticators';
+import {OrganizationFixture} from 'sentry-fixture/organization';
 import {RouterContextFixture} from 'sentry-fixture/routerContextFixture';
 import {RouterFixture} from 'sentry-fixture/routerFixture';
 
 import {render, screen, userEvent} from 'sentry-test/reactTestingLibrary';
 
+import OrganizationsStore from 'sentry/stores/organizationsStore';
 import AccountSecurityEnroll from 'sentry/views/settings/account/accountSecurity/accountSecurityEnroll';
 
 const ENDPOINT = '/users/me/authenticators/';
+const usorg = OrganizationFixture({
+  slug: 'us-org',
+  links: {
+    organizationUrl: 'https://us-org.example.test',
+    regionUrl: 'https://us.example.test',
+  },
+});
 
 describe('AccountSecurityEnroll', function () {
+  jest.spyOn(window.location, 'assign').mockImplementation(() => {});
+
   describe('Totp', function () {
     const authenticator = AuthenticatorsFixture().Totp({
       isEnrolled: false,
@@ -32,14 +43,31 @@ describe('AccountSecurityEnroll', function () {
       },
     ]);
 
+    let location;
     beforeEach(function () {
+      location = window.location;
+      window.location.href = 'https://example.test';
+      window.__initialData = {
+        ...window.__initialData,
+        links: {
+          organizationUrl: undefined,
+          regionUrl: undefined,
+          sentryUrl: 'https://example.test',
+        },
+      };
+      OrganizationsStore.load([usorg]);
+
       MockApiClient.clearMockResponses();
       MockApiClient.addMockResponse({
         url: `${ENDPOINT}${authenticator.authId}/enroll/`,
         body: authenticator,
       });
     });
 
+    beforeEach(function () {
+      window.location = location;
+    });
+
     it('does not have enrolled circle indicator', function () {
       render(<AccountSecurityEnroll />, {context: routerContext});
 
@@ -54,11 +82,63 @@ describe('AccountSecurityEnroll', function () {
       expect(screen.getByLabelText('Enrollment QR Code')).toBeInTheDocument();
     });
 
-    it('can enroll', async function () {
+    it('can enroll from org subdomain', async function () {
+      window.location.href = 'https://us-org.example.test';
+      window.__initialData = {
+        ...window.__initialData,
+        links: {
+          organizationUrl: 'https://us-org.example.test',
+          regionUrl: 'https://us.example.test',
+          sentryUrl: 'https://example.test',
+        },
+      };
+
+      const enrollMock = MockApiClient.addMockResponse({
+        url: `${ENDPOINT}${authenticator.authId}/enroll/`,
+        method: 'POST',
+      });
+      const fetchOrgsMock = MockApiClient.addMockResponse({
+        url: `/organizations/`,
+        body: [usorg],
+      });
+
+      render(<AccountSecurityEnroll />, {context: routerContext});
+
+      await userEvent.type(screen.getByRole('textbox', {name: 'OTP Code'}), 'otp{enter}');
+
+      expect(enrollMock).toHaveBeenCalledWith(
+        `${ENDPOINT}15/enroll/`,
+        expect.objectContaining({
+          method: 'POST',
+          data: expect.objectContaining({
+            secret: 'secret',
+            otp: 'otp',
+          }),
+        })
+      );
+      expect(fetchOrgsMock).not.toHaveBeenCalled();
+      expect(window.location.assign).not.toHaveBeenCalled();
+    });
+
+    it('can enroll from main domain', async function () {
+      OrganizationsStore.load([]);
+      window.__initialData = {
+        ...window.__initialData,
+        links: {
+          organizationUrl: 'https://us-org.example.test',
+          regionUrl: 'https://us.example.test',
+          sentryUrl: 'https://example.test',
+        },
+      };
+
       const enrollMock = MockApiClient.addMockResponse({
         url: `${ENDPOINT}${authenticator.authId}/enroll/`,
         method: 'POST',
       });
+      const fetchOrgsMock = MockApiClient.addMockResponse({
+        url: `/organizations/`,
+        body: [usorg],
+      });
 
       render(<AccountSecurityEnroll />, {context: routerContext});
 
@@ -74,6 +154,9 @@ describe('AccountSecurityEnroll', function () {
           }),
         })
       );
+      expect(fetchOrgsMock).toHaveBeenCalledTimes(1);
+      expect(window.location.assign).toHaveBeenCalledTimes(1);
+      expect(window.location.assign).toHaveBeenCalledWith('http://us-org.example.test/');
     });
 
     it('can redirect with already enrolled error', function () {

static/app/views/settings/account/accountSecurity/accountSecurityEnroll.tsx

@@ -9,7 +9,10 @@ import {
   addSuccessMessage,
 } from 'sentry/actionCreators/indicator';
 import {openRecoveryOptions} from 'sentry/actionCreators/modal';
-import {fetchOrganizationByMember} from 'sentry/actionCreators/organizations';
+import {
+  fetchOrganizationByMember,
+  fetchOrganizations,
+} from 'sentry/actionCreators/organizations';
 import {Alert} from 'sentry/components/alert';
 import {Button} from 'sentry/components/button';
 import ButtonBar from 'sentry/components/buttonBar';
@@ -24,8 +27,10 @@ import PanelItem from 'sentry/components/panels/panelItem';
 import TextCopyInput from 'sentry/components/textCopyInput';
 import U2fsign from 'sentry/components/u2f/u2fsign';
 import {t} from 'sentry/locale';
+import OrganizationsStore from 'sentry/stores/organizationsStore';
 import {space} from 'sentry/styles/space';
 import type {Authenticator} from 'sentry/types';
+import {generateOrgSlugUrl} from 'sentry/utils';
 import getPendingInvite from 'sentry/utils/getPendingInvite';
 // eslint-disable-next-line no-restricted-imports
 import withSentryRouter from 'sentry/utils/withSentryRouter';
@@ -330,6 +335,34 @@ class AccountSecurityEnroll extends DeprecatedAsyncView<Props, State> {
 
     this.props.router.push('/settings/account/security/');
     openRecoveryOptions({authenticatorName: this.authenticatorName});
+
+    // The remainder of this function is included primarily to smooth out the relocation flow. The
+    // newly claimed user will have landed on `https://sentry.io/settings/account/security` to
+    // perform the 2FA registration. But now that they have in fact registered, we want to redirect
+    // them to the subdomain of the organization they are already a member of (ex:
+    // `https://my-2fa-org.sentry.io`), but did not have the ability to access due to their previous
+    // lack of 2FA enrollment.
+    let orgs = OrganizationsStore.getAll();
+    if (orgs.length === 0) {
+      // Try to load orgs post 2FA again.
+      orgs = await fetchOrganizations(this.api, {member: '1'});
+      OrganizationsStore.load(orgs);
+
+      // Still no orgs? Nowhere to redirect the user to, so just stay in place.
+      if (orgs.length === 0) {
+        return;
+      }
+    }
+
+    // If we are already in an org sub-domain, we don't need to do any redirection. If we are not
+    // (this is usually only the case for a newly claimed relocated user), we redirect to the org
+    // slug's subdomain now.
+    const isAlreadyInOrgSubDomain = orgs.some(org => {
+      return org.links.organizationUrl === new URL(window.location.href).origin;
+    });
+    if (!isAlreadyInOrgSubDomain) {
+      window.location.assign(generateOrgSlugUrl(orgs[0].slug));
+    }
   }
 
   // Handler when we failed to add a 2fa device