Fix MFA from tray menu

src-tauri/src/events.rs

@@ -17,6 +17,7 @@ pub enum EventKey {
     DeadConnectionReconnected,
     ApplicationConfigChanged,
     AddInstance,
+    MfaTrigger,
 }
 
 impl From<EventKey> for &'static str {
@@ -32,6 +33,7 @@ impl From<EventKey> for &'static str {
             EventKey::DeadConnectionReconnected => "dead-connection-reconnected",
             EventKey::ApplicationConfigChanged => "application-config-changed",
             EventKey::AddInstance => "add-instance",
+            EventKey::MfaTrigger => "mfa-trigger",
         }
     }
 }

src-tauri/src/tray.rs

@@ -12,6 +12,7 @@ use crate::{
     commands::{all_instances, all_locations, connect, disconnect},
     database::{models::location::Location, DB_POOL},
     error::Error,
+    events::EventKey,
     ConnectionType,
 };
 
@@ -225,11 +226,6 @@ pub async fn configure_tray_icon(app: &AppHandle, theme: AppTrayTheme) -> Result
     }
 }
 
-#[derive(Clone, serde::Serialize)]
-struct MfaPayload<'a> {
-    message: &'a str,
-}
-
 async fn handle_location_tray_menu(id: String, handle: &AppHandle) {
     match id.parse::<i64>() {
         Ok(location_id) => {
@@ -244,18 +240,13 @@ async fn handle_location_tray_menu(id: String, handle: &AppHandle) {
                             disconnect(location_id, ConnectionType::Location, handle.clone()).await;
                     } else {
                         info!("Connect location with ID {id}");
-                        // check is mfa enabled and trigger modal on frontend
+                        // Check if MFA is enabled. If so, trigger modal on frontend.
                         if location.mfa_enabled() {
                             info!(
                                 "MFA enabled for location with ID {:?}, trigger MFA modal",
                                 location.id
                             );
-                            let _ = handle.emit(
-                                "mfa-trigger",
-                                MfaPayload {
-                                    message: "Trigger MFA event",
-                                },
-                            );
+                            let _ = handle.emit(EventKey::MfaTrigger.into(), &location);
                         } else if let Err(err) =
                             connect(location_id, ConnectionType::Location, None, handle.clone())
                                 .await

src/pages/client/ClientPage.tsx

@@ -16,9 +16,11 @@ import { DeadConDroppedModal } from './components/modals/DeadConDroppedModal/Dea
 import { useDeadConDroppedModal } from './components/modals/DeadConDroppedModal/store';
 import { useClientFlags } from './hooks/useClientFlags';
 import { useClientStore } from './hooks/useClientStore';
+import { useMFAModal } from './pages/ClientInstancePage/components/LocationsList/modals/MFAModal/useMFAModal';
 import { clientQueryKeys } from './query';
 import {
   type AddInstancePayload,
+  type CommonWireguardFields,
   type DeadConDroppedPayload,
   type DeadConReconnectedPayload,
   TauriEventKey,
@@ -41,6 +43,7 @@ export const ClientPage = () => {
   const location = useLocation();
   const toaster = useToaster();
   const openDeadConDroppedModal = useDeadConDroppedModal((s) => s.open);
+  const openMFAModal = useMFAModal((state) => state.open);
   const { LL } = useI18nContext();
 
   const { data: instances } = useQuery({
@@ -137,11 +140,15 @@ export const ClientPage = () => {
       },
     );
 
-    const doEnrollment = listen<AddInstancePayload>(
-      TauriEventKey.ADD_INSTANCE,
+    const addInstance = listen<AddInstancePayload>(TauriEventKey.ADD_INSTANCE, (data) => {
+      useClientStore.setState({ instanceConfig: data.payload });
+      navigate(routes.client.addInstance, { replace: true });
+    });
+
+    const mfaTrigger = listen<CommonWireguardFields>(
+      TauriEventKey.MFA_TRIGGER,
       (data) => {
-        useClientStore.setState({ instanceConfig: data.payload });
-        navigate(routes.client.addInstance, { replace: true });
+        openMFAModal(data.payload);
       },
     );
 
@@ -153,7 +160,8 @@ export const ClientPage = () => {
       instanceUpdate.then((cleanup) => cleanup());
       locationUpdate.then((cleanup) => cleanup());
       appConfigChanged.then((cleanup) => cleanup());
-      doEnrollment.then((cleanup) => cleanup());
+      addInstance.then((cleanup) => cleanup());
+      mfaTrigger.then((cleanup) => cleanup());
     };
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, []);

src/pages/client/components/ClientSideBar/components/ClientBarItem/ClientBarItem.tsx

@@ -7,7 +7,7 @@ import { useMatch, useNavigate } from 'react-router-dom';
 import SvgIconConnection from '../../../../../../shared/defguard-ui/components/svg/IconConnection';
 import { routes } from '../../../../../../shared/routes';
 import { useClientStore } from '../../../../hooks/useClientStore';
-import { type SelectedInstance, WireguardInstanceType } from '../../../../types';
+import type { WireguardInstanceType } from '../../../../types';
 
 type Props = {
   itemType: WireguardInstanceType;
@@ -47,26 +47,12 @@ export const ClientBarItem = ({ itemType, itemId, label, active = false }: Props
         className={cn}
         ref={refs.setReference}
         onClick={() => {
-          switch (itemType) {
-            case WireguardInstanceType.DEFGUARD_INSTANCE: {
-              const _selectedInstance: SelectedInstance = {
-                id: itemId,
-                type: WireguardInstanceType.DEFGUARD_INSTANCE,
-              };
-              setClientStore({
-                selectedInstance: _selectedInstance,
-              });
-              break;
-            }
-            case WireguardInstanceType.TUNNEL:
-              setClientStore({
-                selectedInstance: {
-                  id: itemId,
-                  type: WireguardInstanceType.TUNNEL,
-                },
-              });
-              break;
-          }
+          setClientStore({
+            selectedInstance: {
+              id: itemId,
+              type: itemType,
+            },
+          });
           if (!instancePage) {
             navigate(routes.client.instancePage, { replace: true });
           }

src/pages/client/pages/ClientInstancePage/components/LocationsList/modals/MFAModal/MFAModal.tsx

@@ -29,7 +29,6 @@ import { useClientStore } from '../../../../../../hooks/useClientStore';
 import {
   type DefguardInstance,
   LocationMfaType,
-  WireguardInstanceType,
 } from '../../../../../../types';
 import { MfaMobileApprove } from './components/MfaMobileApprove/MfaMobileApprove';
 import { BrowserErrorIcon, BrowserPendingIcon, GoToBrowserIcon } from './Icons';
@@ -76,23 +75,15 @@ export const MFAModal = () => {
   const [startResponse, setStartResponse] = useState<MFAStartResponse>();
 
   const localLL = LL.modals.mfa.authentication;
-  const isOpen = useMFAModal((state) => state.isOpen);
-  const location = useMFAModal((state) => state.instance);
+  const [isOpen, location] = useMFAModal((state) => [state.isOpen, state.instance]);
   const [close, reset] = useMFAModal((state) => [state.close, state.reset], shallow);
-  const [selectedInstanceId, selectedInstanceType] = useClientStore((state) => [
-    state.selectedInstance?.id,
-    state.selectedInstance?.type,
-  ]);
   const instances = useClientStore((state) => state.instances);
   const selectedInstance = useMemo((): DefguardInstance | undefined => {
-    if (
-      !isUndefined(selectedInstanceId) &&
-      selectedInstanceType &&
-      selectedInstanceType === WireguardInstanceType.DEFGUARD_INSTANCE
-    ) {
-      return instances.find((i) => i.id === selectedInstanceId);
+    const instanceId = location?.instance_id;
+    if (!isUndefined(instanceId)) {
+      return instances.find((i) => i.id === instanceId);
     }
-  }, [selectedInstanceId, selectedInstanceType, instances]);
+  }, [location, instances]);
 
   const resetState = () => {
     reset();

src/pages/client/types.ts

@@ -119,4 +119,5 @@ export enum TauriEventKey {
   DEAD_CONNECTION_RECONNECTED = 'dead-connection-reconnected',
   APPLICATION_CONFIG_CHANGED = 'application-config-changed',
   ADD_INSTANCE = 'add-instance',
+  MFA_TRIGGER = 'mfa-trigger',
 }

src/pages/enrollment/hooks/store/useEnrollmentStore.tsx

@@ -1,6 +1,6 @@
 import { pick } from 'lodash-es';
 import { Subject } from 'rxjs';
-import { createJSONStorage, devtools, persist } from 'zustand/middleware';
+import { createJSONStorage, persist } from 'zustand/middleware';
 import { createWithEqualityFn } from 'zustand/traditional';
 
 import type { AdminInfo, UserInfo } from '../../../../shared/hooks/api/types';
@@ -34,36 +34,34 @@ const persistKeys: Array<keyof StoreValues> = [
 ];
 
 export const useEnrollmentStore = createWithEqualityFn<Store>()(
-  devtools(
-    persist(
-      (set, get) => ({
-        ...defaultValues,
-        init: (values) => set({ ...defaultValues, ...values }),
-        setState: (newValues) => set((old) => ({ ...old, ...newValues })),
-        reset: () => set(defaultValues),
-        nextStep: () => {
-          const current = get().step;
-          const max = get().stepsMax;
+  persist(
+    (set, get) => ({
+      ...defaultValues,
+      init: (values) => set({ ...defaultValues, ...values }),
+      setState: (newValues) => set((old) => ({ ...old, ...newValues })),
+      reset: () => set(defaultValues),
+      nextStep: () => {
+        const current = get().step;
+        const max = get().stepsMax;
 
-          if (current < max) {
-            return set({ step: current + 1 });
-          }
-        },
-        perviousStep: () => {
-          const current = get().step;
+        if (current < max) {
+          return set({ step: current + 1 });
+        }
+      },
+      perviousStep: () => {
+        const current = get().step;
 
-          if (current > 0) {
-            return set({ step: current - 1 });
-          }
-        },
-      }),
-      {
-        name: 'enrollment-storage',
-        version: 0.1,
-        storage: createJSONStorage(() => sessionStorage),
-        partialize: (state) => pick(state, persistKeys),
+        if (current > 0) {
+          return set({ step: current - 1 });
+        }
       },
-    ),
+    }),
+    {
+      name: 'enrollment-storage',
+      version: 0.1,
+      storage: createJSONStorage(() => sessionStorage),
+      partialize: (state) => pick(state, persistKeys),
+    },
   ),
   Object.is,
 );