Merge branch 'dev' of github.com:DefGuard/client into stats_wip

src-tauri/src/commands.rs

@@ -179,12 +179,17 @@ pub async fn all_instances(app_state: State<'_, AppState>) -> Result<Vec<Instanc
         let connected = connection_ids
             .iter()
             .any(|item1| location_ids.iter().any(|item2| item1 == item2));
+        let keys = WireguardKeys::find_by_instance_id(&app_state.get_pool(), instance.id.unwrap())
+            .await
+            .map_err(|err| err.to_string())?
+            .unwrap();
         instance_info.push(InstanceInfo {
             id: instance.id,
             uuid: instance.uuid.clone(),
             name: instance.name.clone(),
             url: instance.url.clone(),
             connected,
+            pubkey: keys.pubkey,
         })
     }
     Ok(instance_info)
@@ -199,7 +204,6 @@ pub struct LocationInfo {
     pub address: String,
     pub endpoint: String,
     pub active: bool,
-    pub pubkey: String,
 }
 
 #[tauri::command(async)]
@@ -218,10 +222,6 @@ pub async fn all_locations(
         .map(|con| con.location_id)
         .collect();
     let mut location_info = vec![];
-    let keys = WireguardKeys::find_by_instance_id(&app_state.get_pool(), instance_id)
-        .await
-        .map_err(|err| err.to_string())?
-        .unwrap();
     for location in locations {
         let info = LocationInfo {
             id: location.id.unwrap(),
@@ -230,7 +230,6 @@ pub async fn all_locations(
             address: location.address,
             endpoint: location.endpoint,
             active: active_locations_ids.contains(&location.id.unwrap()),
-            pubkey: keys.pubkey.clone(),
         };
         location_info.push(info);
     }

src-tauri/src/database/models/instance.rs

@@ -50,4 +50,5 @@ pub struct InstanceInfo {
     pub uuid: String,
     pub url: String,
     pub connected: bool,
+    pub pubkey: String,
 }

src/pages/client/components/modals/AddInstanceModal/AddInstanceModal.tsx

@@ -4,6 +4,7 @@ import { shallow } from 'zustand/shallow';
 import { useI18nContext } from '../../../../../i18n/i18n-react';
 import { ModalWithTitle } from '../../../../../shared/defguard-ui/components/Layout/modals/ModalWithTitle/ModalWithTitle';
 import { useAddInstanceModal } from './hooks/useAddInstanceModal';
+import { AddInstanceDeviceStep } from './steps/AddInstanceDeviceStep';
 import { AddInstanceModalInitStep } from './steps/AddInstanceInitStep';
 
 export const AddInstanceModal = () => {
@@ -45,4 +46,7 @@ export const AddInstanceModal = () => {
   );
 };
 
-const steps: ReactNode[] = [<AddInstanceModalInitStep key={0} />];
+const steps: ReactNode[] = [
+  <AddInstanceModalInitStep key={0} />,
+  <AddInstanceDeviceStep key={1} />,
+];

src/pages/client/components/modals/AddInstanceModal/steps/AddInstanceDeviceStep.tsx

@@ -1,4 +1,5 @@
 import { zodResolver } from '@hookform/resolvers/zod';
+import { useQueryClient } from '@tanstack/react-query';
 import { invoke } from '@tauri-apps/api';
 import { useMemo, useState } from 'react';
 import { SubmitHandler, useForm } from 'react-hook-form';
@@ -18,6 +19,7 @@ import {
   CreateDeviceResponse,
 } from '../../../../../../shared/hooks/api/types';
 import { generateWGKeys } from '../../../../../../shared/utils/generateWGKeys';
+import { clientQueryKeys } from '../../../../query';
 import { useAddInstanceModal } from '../hooks/useAddInstanceModal';
 
 type FormFields = {
@@ -34,6 +36,7 @@ export const AddInstanceDeviceStep = () => {
   const toaster = useToaster();
   const close = useAddInstanceModal((state) => state.close);
   const [isLoading, setIsLoading] = useState(false);
+  const queryClient = useQueryClient();
 
   const [proxyUrl] = useAddInstanceModal((state) => [state.proxyUrl], shallow);
 
@@ -78,6 +81,8 @@ export const AddInstanceDeviceStep = () => {
             .then(() => {
               setIsLoading(false);
               toaster.success(componentLL.messages.success.add());
+              queryClient.invalidateQueries([clientQueryKeys.getInstances]);
+              queryClient.invalidateQueries([clientQueryKeys.getLocations]);
               close();
             })
             .catch((e) => {

src/pages/client/components/modals/AddInstanceModal/steps/AddInstanceInitStep.tsx

@@ -1,5 +1,6 @@
 import { zodResolver } from '@hookform/resolvers/zod';
 import dayjs from 'dayjs';
+import { isUndefined } from 'lodash-es';
 import { useMemo } from 'react';
 import { SubmitHandler, useForm } from 'react-hook-form';
 import { useNavigate } from 'react-router-dom';
@@ -16,6 +17,7 @@ import { useToaster } from '../../../../../../shared/defguard-ui/hooks/toasts/us
 import { EnrollmentStartResponse } from '../../../../../../shared/hooks/api/types';
 import { routes } from '../../../../../../shared/routes';
 import { useEnrollmentStore } from '../../../../../enrollment/hooks/store/useEnrollmentStore';
+import { clientApi } from '../../../../clientAPI/clientApi';
 import { useAddInstanceModal } from '../hooks/useAddInstanceModal';
 
 type FormFields = {
@@ -93,13 +95,29 @@ export const AddInstanceModalInitStep = () => {
           setModalState({ loading: false });
           return;
         }
-        res.json().then((r: EnrollmentStartResponse) => {
-          setModalState({ loading: false });
+        res.json().then(async (r: EnrollmentStartResponse) => {
+          // get client registered instances
+          const clientInstances = await clientApi.getInstances();
+          const instance = clientInstances.find((i) => i.uuid === r.instance.id);
           let proxy_api_url = import.meta.env.DEV ? '' : values.url;
           if (proxy_api_url[proxy_api_url.length - 1] === '/') {
             proxy_api_url = proxy_api_url.slice(0, -1);
           }
           proxy_api_url = proxy_api_url + '/api/v1';
+          setModalState({ loading: false });
+
+          if (instance) {
+            // update already registered instance instead
+            const instanceInfo = await fetch(`${proxy_api_url}/enrollment/network_info`, {
+              method: 'POST',
+              headers,
+              body: JSON.stringify({
+                pubKey: instance.pubkey,
+              }),
+            });
+            return;
+          }
+          // register new instance
           // is user in need of full enrollment ?
           if (r.user.is_active) {
             //no, only create new device for desktop client

src/pages/client/types.ts

@@ -4,6 +4,7 @@ export type DefguardInstance = {
   name: string;
   url: string;
   connected: boolean;
+  pubkey: string;
 };
 
 export type DefguardLocation = {
@@ -12,7 +13,6 @@ export type DefguardLocation = {
   network_id: number;
   name: string;
   address: string;
-  pubkey: string;
   endpoint: string;
   allowed_ips: string;
 };