refactor: improve workflow

gpgui/src/atoms/gateway.ts

@@ -22,7 +22,8 @@ export const gatewayLoginAtom = atom(
       throw new Error("Failed to login to gateway");
     }
 
-    if (!get(isProcessingAtom)) {
+    const isProcessing = get(isProcessingAtom);
+    if (!isProcessing) {
       console.info("Request cancelled");
       return;
     }
@@ -44,13 +45,21 @@ const connectVpnAtom = atom(
   }
 );
 
+const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
+
 export const disconnectVpnAtom = atom(null, async (get, set) => {
   try {
     set(statusAtom, "disconnecting");
     await vpnService.disconnect();
-    set(statusAtom, "disconnected");
+    // Sleep a short time, so that the client can receive the service's disconnected event.
+    await sleep(100);
   } catch (err) {
     set(statusAtom, "disconnected");
     set(notifyErrorAtom, "Failed to disconnect from VPN");
   }
 });
+
+export const gatewaySwitcherVisibleAtom = atom(false);
+export const openGatewaySwitcherAtom = atom(null, (get, set) => {
+  set(gatewaySwitcherVisibleAtom, true);
+});

gpgui/src/atoms/menu.ts

@@ -0,0 +1,20 @@
+import { exit } from "@tauri-apps/api/process";
+import { atom } from "jotai";
+import { RESET } from "jotai/utils";
+import { disconnectVpnAtom } from "./gateway";
+import { appDataStorageAtom, portalAddressAtom } from "./portal";
+import { statusAtom } from "./status";
+
+export const resetAtom = atom(null, (_get, set) => {
+  set(appDataStorageAtom, RESET);
+  set(portalAddressAtom, "");
+});
+
+export const quitAtom = atom(null, async (get, set) => {
+  const status = get(statusAtom);
+
+  if (status === "connected") {
+    await set(disconnectVpnAtom);
+  }
+  await exit();
+});

gpgui/src/atoms/notification.ts

@@ -3,11 +3,19 @@ import { atom } from "jotai";
 
 export type Severity = AlertColor;
 
+type NotificationConfig = {
+  title: string;
+  message: string;
+  severity: Severity;
+  duration?: number;
+};
+
 const notificationVisibleAtom = atom(false);
-export const notificationConfigAtom = atom({
+export const notificationConfigAtom = atom<NotificationConfig>({
   title: "",
   message: "",
   severity: "info" as Severity,
+  duration: 5000,
 });
 
 export const closeNotificationAtom = atom(
@@ -17,20 +25,37 @@ export const closeNotificationAtom = atom(
   }
 );
 
-export const notifyErrorAtom = atom(null, (_get, set, err: unknown) => {
-  let msg: string;
-  if (err instanceof Error) {
-    msg = err.message;
-  } else if (typeof err === "string") {
-    msg = err;
-  } else {
-    msg = "Unknown error";
-  }
+export const notifyErrorAtom = atom(
+  null,
+  (_get, set, err: unknown, duration: number = 5000) => {
+    let msg: string;
+    if (err instanceof Error) {
+      msg = err.message;
+    } else if (typeof err === "string") {
+      msg = err;
+    } else {
+      msg = "Unknown error";
+    }
 
-  set(notificationVisibleAtom, true);
-  set(notificationConfigAtom, {
-    title: "Error",
-    message: msg,
-    severity: "error",
-  });
-});
+    set(notificationVisibleAtom, true);
+    set(notificationConfigAtom, {
+      title: "Error",
+      message: msg,
+      severity: "error",
+      duration: duration <= 0 ? undefined : duration,
+    });
+  }
+);
+
+export const notifySuccessAtom = atom(
+  null,
+  (_get, set, msg: string, duration: number = 5000) => {
+    set(notificationVisibleAtom, true);
+    set(notificationConfigAtom, {
+      title: "Success",
+      message: msg,
+      severity: "success",
+      duration: duration <= 0 ? undefined : duration,
+    });
+  }
+);

gpgui/src/atoms/portal.ts

@@ -1,42 +1,119 @@
 import { atom } from "jotai";
-import { focusAtom } from "jotai-optics";
+import { withImmer } from "jotai-immer";
+import { atomWithDefault, atomWithStorage } from "jotai/utils";
 import authService, { AuthData } from "../services/authService";
 import portalService, {
   PasswordPrelogin,
+  PortalCredential,
   Prelogin,
   SamlPrelogin,
 } from "../services/portalService";
-import { gatewayLoginAtom } from "./gateway";
+import { disconnectVpnAtom, gatewayLoginAtom } from "./gateway";
 import { notifyErrorAtom } from "./notification";
 import { isProcessingAtom, statusAtom } from "./status";
 
-type GatewayData = {
+export type GatewayData = {
   name: string;
   address: string;
 };
 
-type Credential = {
-  user: string;
-  passwd: string;
-  userAuthCookie: string;
-  prelogonUserAuthCookie: string;
+type CachedPortalCredential = Omit<PortalCredential, "prelogin-cookie">;
+
+type PortalData = {
+  address: string;
+  gateways: GatewayData[];
+  cachedCredential?: CachedPortalCredential;
+  selectedGateway?: string;
 };
 
 type AppData = {
   portal: string;
-  gateways: GatewayData[];
-  selectedGateway: string;
-  credentials: Record<string, Credential>;
+  portals: PortalData[];
+  clearCookies: boolean;
 };
 
-const appAtom = atom<AppData>({
+type AppDataUpdate =
+  | {
+      type: "PORTAL";
+      payload: PortalData;
+    }
+  | {
+      type: "SELECTED_GATEWAY";
+      payload: string;
+    };
+
+const defaultAppData: AppData = {
   portal: "",
-  gateways: [],
-  selectedGateway: "",
-  credentials: {},
+  portals: [],
+  // Whether to clear the cookies of the SAML login webview, default is true
+  clearCookies: true,
+};
+
+export const appDataStorageAtom = atomWithStorage<AppData>(
+  "APP_DATA",
+  defaultAppData
+);
+const appDataImmerAtom = withImmer(appDataStorageAtom);
+
+const updateAppDataAtom = atom(null, (_get, set, update: AppDataUpdate) => {
+  const { type, payload } = update;
+  switch (type) {
+    case "PORTAL":
+      const { address } = payload;
+      set(appDataImmerAtom, (draft) => {
+        draft.portal = address;
+        const portalIndex = draft.portals.findIndex(
+          ({ address: portalAddress }) => portalAddress === address
+        );
+        if (portalIndex === -1) {
+          draft.portals.push(payload);
+        } else {
+          draft.portals[portalIndex] = payload;
+        }
+      });
+      break;
+    case "SELECTED_GATEWAY":
+      set(appDataImmerAtom, (draft) => {
+        const { portal, portals } = draft;
+        const portalData = portals.find(({ address }) => address === portal);
+        if (portalData) {
+          portalData.selectedGateway = payload;
+        }
+      });
+      break;
+  }
 });
 
-export const portalAtom = focusAtom(appAtom, (optic) => optic.prop("portal"));
+export const portalAddressAtom = atomWithDefault(
+  (get) => get(appDataImmerAtom).portal
+);
+
+export const currentPortalDataAtom = atom<PortalData>((get) => {
+  const portalAddress = get(portalAddressAtom);
+  const { portals } = get(appDataImmerAtom);
+  const portalData = portals.find(({ address }) => address === portalAddress);
+
+  return portalData || { address: portalAddress, gateways: [] };
+});
+
+const clearCookiesAtom = atom(
+  (get) => get(appDataImmerAtom).clearCookies,
+  (_get, set, update: boolean) => {
+    set(appDataImmerAtom, (draft) => {
+      draft.clearCookies = update;
+    });
+  }
+);
+
+export const portalGatewaysAtom = atom<GatewayData[]>((get) => {
+  const { gateways } = get(currentPortalDataAtom);
+  return gateways;
+});
+
+export const selectedGatewayAtom = atom(
+  (get) => get(currentPortalDataAtom).selectedGateway
+);
+
 export const connectPortalAtom = atom(
   (get) => get(isProcessingAtom),
   async (get, set, action?: "retry-auth") => {
@@ -46,7 +123,7 @@ export const connectPortalAtom = atom(
       return;
     }
 
-    const portal = get(portalAtom);
+    const portal = get(portalAddressAtom);
     if (!portal) {
       set(notifyErrorAtom, "Portal is empty");
       return;
@@ -55,15 +132,20 @@ export const connectPortalAtom = atom(
     try {
       set(statusAtom, "prelogin");
       const prelogin = await portalService.prelogin(portal);
-      if (!get(isProcessingAtom)) {
+      const isProcessing = get(isProcessingAtom);
+      if (!isProcessing) {
         console.info("Request cancelled");
         return;
       }
 
-      if (prelogin.isSamlAuth) {
-        await set(launchSamlAuthAtom, prelogin);
-      } else {
-        await set(launchPasswordAuthAtom, prelogin);
+      try {
+        await set(loginWithCachedCredentialAtom, prelogin);
+      } catch {
+        if (prelogin.isSamlAuth) {
+          await set(launchSamlAuthAtom, prelogin);
+        } else {
+          await set(launchPasswordAuthAtom, prelogin);
+        }
       }
     } catch (err) {
       set(cancelConnectPortalAtom);
@@ -78,6 +160,17 @@ connectPortalAtom.onMount = (dispatch) => {
   });
 };
 
+const loginWithCachedCredentialAtom = atom(
+  null,
+  async (get, set, prelogin: Prelogin) => {
+    const { cachedCredential } = get(currentPortalDataAtom);
+    if (!cachedCredential) {
+      throw new Error("No cached credential");
+    }
+    await set(portalLoginAtom, cachedCredential, prelogin);
+  }
+);
+
 export const passwordPreloginAtom = atom<PasswordPrelogin>({
   isSamlAuth: false,
   region: "",
@@ -90,8 +183,14 @@ export const cancelConnectPortalAtom = atom(null, (_get, set) => {
   set(statusAtom, "disconnected");
 });
 
-export const usernameAtom = atom("");
-export const passwordAtom = atom("");
+export const usernameAtom = atomWithDefault(
+  (get) => get(currentPortalDataAtom).cachedCredential?.user ?? ""
+);
+
+export const passwordAtom = atomWithDefault(
+  (get) => get(currentPortalDataAtom).cachedCredential?.passwd ?? ""
+);
+
 const passwordAuthVisibleAtom = atom(false);
 
 const launchPasswordAuthAtom = atom(
@@ -114,7 +213,7 @@ export const cancelPasswordAuthAtom = atom(
 export const passwordLoginAtom = atom(
   (get) => get(portalConfigLoadingAtom),
   async (get, set, username: string, password: string) => {
-    const portal = get(portalAtom);
+    const portal = get(portalAddressAtom);
     if (!portal) {
       set(notifyErrorAtom, "Portal is empty");
       return;
@@ -138,13 +237,18 @@ export const passwordLoginAtom = atom(
 
 const launchSamlAuthAtom = atom(
   null,
-  async (_get, set, prelogin: SamlPrelogin) => {
+  async (get, set, prelogin: SamlPrelogin) => {
     const { samlAuthMethod, samlRequest } = prelogin;
     let authData: AuthData;
 
     try {
       set(statusAtom, "authenticating-saml");
-      authData = await authService.samlLogin(samlAuthMethod, samlRequest);
+      const clearCookies = get(clearCookiesAtom);
+      authData = await authService.samlLogin(
+        samlAuthMethod,
+        samlRequest,
+        clearCookies
+      );
     } catch (err) {
       throw new Error("SAML login failed");
     }
@@ -155,17 +259,21 @@ const launchSamlAuthAtom = atom(
       return;
     }
 
+    // SAML login success, update clearCookies to false to reuse the SAML session
+    set(clearCookiesAtom, false);
+
     const credential = {
       user: authData.username,
       "prelogin-cookie": authData.prelogin_cookie,
       "portal-userauthcookie": authData.portal_userauthcookie,
     };
+
     await set(portalLoginAtom, credential, prelogin);
   }
 );
 
 const retrySamlAuthAtom = atom(null, async (get) => {
-  const portal = get(portalAtom);
+  const portal = get(portalAddressAtom);
   const prelogin = await portalService.prelogin(portal);
   if (prelogin.isSamlAuth) {
     await authService.emitAuthRequest({
@@ -175,17 +283,6 @@ const retrySamlAuthAtom = atom(null, async (get) => {
   }
 });
 
-type PortalCredential =
-  | {
-      user: string;
-      passwd: string;
-    }
-  | {
-      user: string;
-      "prelogin-cookie": string | null;
-      "portal-userauthcookie": string | null;
-    };
-
 const portalConfigLoadingAtom = atom(false);
 const portalLoginAtom = atom(
   (get) => get(portalConfigLoadingAtom),
@@ -193,33 +290,88 @@ const portalLoginAtom = atom(
     set(statusAtom, "portal-config");
     set(portalConfigLoadingAtom, true);
 
-    const portal = get(portalAtom);
+    const portalAddress = get(portalAddressAtom);
     let portalConfig;
     try {
-      portalConfig = await portalService.fetchConfig(portal, credential);
+      portalConfig = await portalService.fetchConfig(portalAddress, credential);
       // Ensure the password auth window is closed
       set(passwordAuthVisibleAtom, false);
     } finally {
       set(portalConfigLoadingAtom, false);
     }
 
-    if (!get(isProcessingAtom)) {
+    const isProcessing = get(isProcessingAtom);
+    if (!isProcessing) {
       console.info("Request cancelled");
       return;
     }
 
     const { gateways, userAuthCookie, prelogonUserAuthCookie } = portalConfig;
-    console.info("portalConfig", portalConfig);
     if (!gateways.length) {
       throw new Error("No gateway found");
     }
 
+    if (userAuthCookie === "empty" || prelogonUserAuthCookie === "empty") {
+      throw new Error("Failed to login, please try again");
+    }
+
+    // Previous selected gateway
+    const previousGateway = get(selectedGatewayAtom);
+    // Update the app data to persist the portal data
+    set(updateAppDataAtom, {
+      type: "PORTAL",
+      payload: {
+        address: portalAddress,
+        gateways: gateways.map(({ name, address }) => ({
+          name,
+          address,
+        })),
+        cachedCredential: {
+          user: credential.user,
+          passwd: credential.passwd,
+          "portal-userauthcookie": userAuthCookie,
+          "portal-prelogonuserauthcookie": prelogonUserAuthCookie,
+        },
+        selectedGateway: previousGateway,
+      },
+    });
+
     const { region } = prelogin;
-    const { address } = portalService.preferredGateway(gateways, region);
+    const { name, address } = portalService.preferredGateway(gateways, {
+      region,
+      previousGateway,
+    });
     await set(gatewayLoginAtom, address, {
       user: credential.user,
       userAuthCookie,
       prelogonUserAuthCookie,
     });
+
+    // Update the app data to persist the gateway data
+    set(updateAppDataAtom, {
+      type: "SELECTED_GATEWAY",
+      payload: name,
+    });
+  }
+);
+
+export const switchingGatewayAtom = atom(false);
+export const switchToGatewayAtom = atom(
+  (get) => get(switchingGatewayAtom),
+  async (get, set, gateway: GatewayData) => {
+    set(updateAppDataAtom, {
+      type: "SELECTED_GATEWAY",
+      payload: gateway.name,
+    });
+
+    if (get(statusAtom) === "connected") {
+      try {
+        set(switchingGatewayAtom, true);
+        await set(disconnectVpnAtom);
+        await set(connectPortalAtom);
+      } finally {
+        set(switchingGatewayAtom, false);
+      }
+    }
   }
 );

gpgui/src/atoms/status.ts

@@ -1,5 +1,8 @@
 import { atom } from "jotai";
+import { atomWithDefault } from "jotai/utils";
 import vpnService from "../services/vpnService";
+import { notifyErrorAtom, notifySuccessAtom } from "./notification";
+import { selectedGatewayAtom, switchingGatewayAtom } from "./portal";
 
 export type Status =
   | "disconnected"
@@ -13,13 +16,42 @@ export type Status =
   | "disconnecting"
   | "error";
 
-export const statusAtom = atom<Status>("disconnected");
-statusAtom.onMount = (setAtom) => {
-  return vpnService.onStatusChanged((status) => {
-    status === "connected" && setAtom("connected");
-  });
+const internalIsOnlineAtom = atomWithDefault(() => vpnService.isOnline());
+export const isOnlineAtom = atom(
+  (get) => get(internalIsOnlineAtom),
+  async (get, set, update: boolean) => {
+    const isOnline = await get(internalIsOnlineAtom);
+    // Already online, do nothing
+    if (update && update === isOnline) {
+      return;
+    }
+
+    set(internalIsOnlineAtom, update);
+    if (update) {
+      set(notifySuccessAtom, "The background service is online");
+    } else {
+      set(notifyErrorAtom, "The background service is offline", 0);
+    }
+  }
+);
+isOnlineAtom.onMount = (setAtom) => vpnService.onServiceStatusChanged(setAtom);
+
+const internalStatusReadyAtom = atom(false);
+export const statusReadyAtom = atom(
+  (get) => get(internalStatusReadyAtom),
+  (get, set, status: Status) => {
+    set(internalStatusReadyAtom, true);
+    set(statusAtom, status);
+  }
+);
+
+statusReadyAtom.onMount = (setAtom) => {
+  vpnService.status().then(setAtom);
 };
 
+export const statusAtom = atom<Status>("disconnected");
+statusAtom.onMount = (setAtom) => vpnService.onVpnStatusChanged(setAtom);
+
 const statusTextMap: Record<Status, String> = {
   disconnected: "Not Connected",
   prelogin: "Portal pre-logging in...",
@@ -35,10 +67,28 @@ const statusTextMap: Record<Status, String> = {
 
 export const statusTextAtom = atom((get) => {
   const status = get(statusAtom);
+  const switchingGateway = get(switchingGatewayAtom);
+
+  if (status === "connected") {
+    const selectedGateway = get(selectedGatewayAtom);
+    return selectedGateway
+      ? `Gateway: ${selectedGateway}`
+      : statusTextMap[status];
+  }
+
+  if (switchingGateway) {
+    const selectedGateway = get(selectedGatewayAtom);
+    return `Switching to ${selectedGateway}`;
+  }
+
   return statusTextMap[status];
 });
 
 export const isProcessingAtom = atom((get) => {
   const status = get(statusAtom);
-  return status !== "disconnected" && status !== "connected";
+  const switchingGateway = get(switchingGatewayAtom);
+
+  return (
+    (status !== "disconnected" && status !== "connected") || switchingGateway
+  );
 });