feat: add cookie import

permitio · Jan 21, 2025 · 6819eb5 · 6819eb5
1 parent b73a953
commit 6819eb5
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 15 deletions.
diff --git a/source/components/AuthProvider.tsx b/source/components/AuthProvider.tsx
@@ -14,6 +14,7 @@ import { ActiveState } from './EnvironmentSelection.js';
 import LoginFlow from './LoginFlow.js';
 import SelectOrganization from './SelectOrganization.js';
 import SelectProject from './SelectProject.js';
+import { useAuthApi } from '../hooks/useAuthApi.js';
 
 // Define the AuthContext type
 type AuthContextType = {
@@ -46,6 +47,7 @@ export function AuthProvider({
 	);
 	const [authToken, setAuthToken] = useState<string>('');
 	const [cookie, setCookie] = useState<string | null>(null);
+	const [newCookie, setNewCookie] = useState<string | null>(null);
 	const [error, setError] = useState<string | null>(null);
 	const [state, setState] = useState<
 		'loading' | 'validate' | 'organization' | 'project' | 'login' | 'done'
@@ -55,6 +57,8 @@ export function AuthProvider({
 	const [loading, setLoading] = useState<boolean>(true);
 	const [currentScope, setCurrentScope] = useState<ApiKeyScope | null>(null);
 
+	const { authSwitchOrgs } = useAuthApi();
+
 	useEffect(() => {
 		if (error) {
 			process.exit(1);
@@ -130,6 +134,25 @@ export function AuthProvider({
 		}
 	}, [key, scope, state, validateApiKeyScope]);
 
+	const switchActiveOrganization = useCallback(
+		async (organization_id: string) => {
+			const { headers, error } = await authSwitchOrgs(
+				organization_id,
+				internalAuthToken,
+				cookie,
+			);
+
+			if (error) {
+				setError(`Error while selecting active workspace: ${error}`);
+				return;
+			}
+
+			let newCookie = headers.getSetCookie()[0] ?? '';
+			setNewCookie(newCookie);
+		},
+		[authSwitchOrgs, internalAuthToken, cookie],
+	);
+
 	useEffect(() => {
 		if (
 			(state === 'organization' && organization) ||
@@ -139,21 +162,25 @@ export function AuthProvider({
 				const { response, error } = await getApiKeyList(
 					state === 'organization' ? 'org' : 'project',
 					internalAuthToken ?? '',
-					cookie,
+					newCookie,
 					project,
 				);
 				if (error || response.data.length === 0) {
-					setError(error ?? 'No API Key found');
+					setError(
+						error
+							? `Error while getting api key list ${error}`
+							: 'No API Key found',
+					);
 					return;
 				}
 				const apiKeyId = response.data[0]?.id ?? '';
 				const { response: secret, error: err } = await getApiKeyById(
 					apiKeyId,
 					internalAuthToken ?? '',
-					cookie,
+					newCookie,
 				);
 				if (err) {
-					setError(err);
+					setError(`Error while getting api key by id: ${err}`);
 					return;
 				}
 				setAuthToken(secret.secret ?? '');
@@ -165,7 +192,7 @@ export function AuthProvider({
 			})();
 		}
 	}, [
-		cookie,
+		newCookie,
 		getApiKeyById,
 		getApiKeyList,
 		internalAuthToken,
@@ -205,20 +232,23 @@ export function AuthProvider({
 					{internalAuthToken && cookie && !organization && (
 						<SelectOrganization
 							accessToken={internalAuthToken}
-							onComplete={organization => setOrganization(organization.value)}
+							onComplete={async organization => {
+								setOrganization(organization.value);
+								await switchActiveOrganization(organization.value);
+							}}
 							onError={setError}
 							cookie={cookie}
 						/>
 					)}
 					{state === 'project' &&
 						internalAuthToken &&
-						cookie &&
+						newCookie &&
 						organization &&
 						!project && (
 							<SelectProject
 								accessToken={internalAuthToken}
 								onComplete={project => setProject(project.value)}
-								cookie={cookie}
+								cookie={newCookie}
 								onError={setError}
 							/>
 						)}

diff --git a/tests/components/AuthProvider.test.tsx b/tests/components/AuthProvider.test.tsx
@@ -110,6 +110,14 @@ describe('AuthProvider', () => {
 					{ id: 'org2', name: 'Organization 2' },
 				]
 			)
+		}).mockResolvedValueOnce({
+			ok: true,
+			status: 200,
+			headers: {
+				getSetCookie: () => ['cookie_value'],
+			},
+			json: async () => ({}),
+			error: null,
 		}).mockResolvedValueOnce({
 			ok: true,
 			status: 200,
@@ -171,7 +179,8 @@ describe('AuthProvider', () => {
 			},
 			json: async () => ({}),
 			error: null,
-		}).mockResolvedValueOnce({
+		})
+			.mockResolvedValueOnce({
 			ok: true,
 			status: 200,
 			error: null,
@@ -180,15 +189,17 @@ describe('AuthProvider', () => {
 					{ id: 'org2', name: 'Organization 2' },
 				]
 			)
-		}).mockResolvedValueOnce({
+		})
+			.mockResolvedValueOnce({
 			ok: true,
 			status: 200,
 			headers: {
 				getSetCookie: () => ['cookie_value'],
 			},
 			json: async () => ({}),
 			error: null,
-		}).mockResolvedValueOnce({
+		})
+			.mockResolvedValueOnce({
 			ok: true,
 			status: 200,
 			error: null,
@@ -197,7 +208,8 @@ describe('AuthProvider', () => {
 					{ id: 'proj2', name: 'Project 2' },
 				]
 			)
-		}).mockResolvedValueOnce({
+		})
+			.mockResolvedValueOnce({
 			ok: true,
 			status: 200,
 			error: null,
@@ -206,15 +218,17 @@ describe('AuthProvider', () => {
 					{ id: 'env2', name: 'Env 2' },
 				]
 			)
-		}).mockResolvedValueOnce({
+		})
+			.mockResolvedValueOnce({
 			ok: true,
 			status: 200,
 			error: null,
 			json: async () => ({
 					secret: 'super-secret'
 				}
 			)
-		}).mockResolvedValueOnce({
+		})
+			.mockResolvedValueOnce({
 			ok: true,
 			status: 200,
 			error: null,
@@ -224,7 +238,8 @@ describe('AuthProvider', () => {
 					]
 				}
 			)
-		}).mockResolvedValueOnce({
+		})
+			.mockResolvedValueOnce({
 			ok: true,
 			status: 200,
 			error: null,