From 7c098529f727fe4c277617b68893a69afdb78e50 Mon Sep 17 00:00:00 2001
From: Akhil Mohan <akhilmhdh@gmail.com>
Date: Wed, 21 Feb 2024 14:09:28 +0530
Subject: [PATCH] fix(admin): resolved undefined on redirect after admin signup

---
 backend/src/server/routes/v1/admin-router.ts    |  6 ++++--
 .../services/super-admin/super-admin-service.ts |  8 ++++++--
 frontend/src/hooks/api/admin/mutation.ts        |  8 +++++++-
 .../src/views/admin/SignUpPage/SignUpPage.tsx   | 17 ++++++++++++++---
 4 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/backend/src/server/routes/v1/admin-router.ts b/backend/src/server/routes/v1/admin-router.ts
index da13d5e00..ab069a2dd 100644
--- a/backend/src/server/routes/v1/admin-router.ts
+++ b/backend/src/server/routes/v1/admin-router.ts
@@ -1,6 +1,6 @@
 import { z } from "zod";
 
-import { SuperAdminSchema, UsersSchema } from "@app/db/schemas";
+import { OrganizationsSchema, SuperAdminSchema, UsersSchema } from "@app/db/schemas";
 import { getConfig } from "@app/lib/config/env";
 import { UnauthorizedError } from "@app/lib/errors";
 import { verifySuperAdmin } from "@app/server/plugins/auth/superAdmin";
@@ -72,6 +72,7 @@ export const registerAdminRouter = async (server: FastifyZodProvider) => {
         200: z.object({
           message: z.string(),
           user: UsersSchema,
+          organization: OrganizationsSchema,
           token: z.string(),
           new: z.string()
         })
@@ -82,7 +83,7 @@ export const registerAdminRouter = async (server: FastifyZodProvider) => {
       const serverCfg = await getServerCfg();
       if (serverCfg.initialized)
         throw new UnauthorizedError({ name: "Admin sign up", message: "Admin has been created" });
-      const { user, token } = await server.services.superAdmin.adminSignUp({
+      const { user, token, organization } = await server.services.superAdmin.adminSignUp({
         ...req.body,
         ip: req.realIp,
         userAgent: req.headers["user-agent"] || ""
@@ -109,6 +110,7 @@ export const registerAdminRouter = async (server: FastifyZodProvider) => {
         message: "Successfully set up admin account",
         user: user.user,
         token: token.access,
+        organization,
         new: "123"
       };
     }
diff --git a/backend/src/services/super-admin/super-admin-service.ts b/backend/src/services/super-admin/super-admin-service.ts
index 1144bd414..9cdc55c8a 100644
--- a/backend/src/services/super-admin/super-admin-service.ts
+++ b/backend/src/services/super-admin/super-admin-service.ts
@@ -96,7 +96,11 @@ export const superAdminServiceFactory = ({
 
     const initialOrganizationName = appCfg.INITIAL_ORGANIZATION_NAME ?? "Admin Org";
 
-    await orgService.createOrganization(userInfo.user.id, userInfo.user.email, initialOrganizationName);
+    const organization = await orgService.createOrganization(
+      userInfo.user.id,
+      userInfo.user.email,
+      initialOrganizationName
+    );
 
     await updateServerCfg({ initialized: true });
     const token = await authService.generateUserTokens({
@@ -106,7 +110,7 @@ export const superAdminServiceFactory = ({
       organizationId: undefined
     });
     // TODO(akhilmhdh-pg): telemetry service
-    return { token, user: userInfo };
+    return { token, user: userInfo, organization };
   };
 
   return {
diff --git a/frontend/src/hooks/api/admin/mutation.ts b/frontend/src/hooks/api/admin/mutation.ts
index ea53ff1db..6d25944ef 100644
--- a/frontend/src/hooks/api/admin/mutation.ts
+++ b/frontend/src/hooks/api/admin/mutation.ts
@@ -2,6 +2,7 @@ import { useMutation, useQueryClient } from "@tanstack/react-query";
 
 import { apiRequest } from "@app/config/request";
 
+import { organizationKeys } from "../organization/queries";
 import { User } from "../users/types";
 import { adminQueryKeys } from "./queries";
 import { TCreateAdminUserDTO, TServerConfig } from "./types";
@@ -9,7 +10,11 @@ import { TCreateAdminUserDTO, TServerConfig } from "./types";
 export const useCreateAdminUser = () => {
   const queryClient = useQueryClient();
 
-  return useMutation<{ user: User; token: string }, {}, TCreateAdminUserDTO>({
+  return useMutation<
+    { user: User; token: string; organization: { id: string } },
+    {},
+    TCreateAdminUserDTO
+  >({
     mutationFn: async (opt) => {
       const { data } = await apiRequest.post("/api/v1/admin/signup", opt);
       return data;
@@ -34,6 +39,7 @@ export const useUpdateServerConfig = () => {
     onSuccess: (data) => {
       queryClient.setQueryData(adminQueryKeys.serverConfig(), data);
       queryClient.invalidateQueries(adminQueryKeys.serverConfig());
+      queryClient.invalidateQueries(organizationKeys.getUserOrganizations);
     }
   });
 };
diff --git a/frontend/src/views/admin/SignUpPage/SignUpPage.tsx b/frontend/src/views/admin/SignUpPage/SignUpPage.tsx
index f452e114e..8a26030ce 100644
--- a/frontend/src/views/admin/SignUpPage/SignUpPage.tsx
+++ b/frontend/src/views/admin/SignUpPage/SignUpPage.tsx
@@ -84,6 +84,10 @@ export const SignUpPage = () => {
         tag: userPass.encryptedPrivateKeyTag,
         privateKey
       });
+      // TODO(akhilmhdh): This is such a confusing pattern and too unreliable
+      // Will be refactored in next iteration to make it url based rather than local storage ones
+      // Part of migration to nextjs 14
+      localStorage.setItem("orgData.id", res.organization.id);
       setStep(SignupSteps.BackupKey);
     } catch (err) {
       console.log(err);
@@ -130,8 +134,8 @@ export const SignUpPage = () => {
           >
             <div className="flex flex-col items-center space-y-2 text-center">
               <img src="/images/gradientLogo.svg" height={90} width={120} alt="Infisical logo" />
-              <div className="text-4xl pt-4">Welcome to Infisical</div>
-              <div className="text-bunker-300 pb-4">Create your first Super Admin Account</div>
+              <div className="pt-4 text-4xl">Welcome to Infisical</div>
+              <div className="pb-4 text-bunker-300">Create your first Super Admin Account</div>
             </div>
             <form onSubmit={handleSubmit(handleFormSubmit)}>
               <div className="mt-8">
@@ -199,7 +203,14 @@ export const SignUpPage = () => {
                   )}
                 />
               </div>
-              <Button type="submit" colorSchema="primary" variant="outline_bg" isFullWidth className="mt-4" isLoading={isSubmitting}>
+              <Button
+                type="submit"
+                colorSchema="primary"
+                variant="outline_bg"
+                isFullWidth
+                className="mt-4"
+                isLoading={isSubmitting}
+              >
                 Continue
               </Button>
             </form>