From 303e086df3037542f59fc823ea1161885bc087d2 Mon Sep 17 00:00:00 2001
From: Gao Sun <gao@silverhand.io>
Date: Wed, 1 Mar 2023 23:46:30 +0800
Subject: [PATCH] refactor: optimize tenant requests

---
 packages/console/src/App.tsx                  |  9 ++-
 .../console/src/cloud/pages/Main/Tenants.tsx  | 17 +++--
 .../console/src/cloud/pages/Main/index.tsx    |  8 +--
 .../src/components/AppLoading/Offline.tsx     |  3 +
 .../src/contexts/AppEndpointsProvider.tsx     | 15 ++--
 .../console/src/contexts/TenantsProvider.tsx  | 18 ++++-
 packages/console/src/hooks/use-api.ts         | 10 ++-
 packages/core/src/middleware/koa-body-etag.ts | 29 ++++++++
 packages/core/src/oidc/init.ts                |  2 +
 packages/core/src/routes/well-known.ts        | 69 ++++++++-----------
 10 files changed, 116 insertions(+), 64 deletions(-)
 create mode 100644 packages/core/src/middleware/koa-body-etag.ts

diff --git a/packages/console/src/App.tsx b/packages/console/src/App.tsx
index 9d094b01e..b3261f093 100644
--- a/packages/console/src/App.tsx
+++ b/packages/console/src/App.tsx
@@ -10,11 +10,12 @@ import './scss/overlayscrollbars.scss';
 
 // eslint-disable-next-line import/no-unassigned-import
 import '@fontsource/roboto-mono';
+
 import CloudApp from '@/cloud/App';
 import { cloudApi, getManagementApi, meApi } from '@/consts/resources';
 import initI18n from '@/i18n/init';
 
-import { adminTenantEndpoint, getUserTenantId } from './consts';
+import { adminTenantEndpoint } from './consts';
 import { isCloud } from './consts/cloud';
 import AppConfirmModalProvider from './contexts/AppConfirmModalProvider';
 import AppEndpointsProvider from './contexts/AppEndpointsProvider';
@@ -24,10 +25,12 @@ import Main from './pages/Main';
 void initI18n();
 
 const Content = () => {
-  const { tenants, isSettle } = useContext(TenantsContext);
-  const currentTenantId = getUserTenantId();
+  const { tenants, isSettle, currentTenantId } = useContext(TenantsContext);
 
   const resources = deduplicate([
+    // Explicitly add `currentTenantId` and deduplicate since the user may directly
+    // access a URL with Tenant ID, adding the ID from the URL here can possibly remove one
+    // additional redirect.
     ...(currentTenantId && [getManagementApi(currentTenantId).indicator]),
     ...(tenants ?? []).map(({ id }) => getManagementApi(id).indicator),
     ...(isCloud ? [cloudApi.indicator] : []),
diff --git a/packages/console/src/cloud/pages/Main/Tenants.tsx b/packages/console/src/cloud/pages/Main/Tenants.tsx
index 67eaa69f7..421b7dac4 100644
--- a/packages/console/src/cloud/pages/Main/Tenants.tsx
+++ b/packages/console/src/cloud/pages/Main/Tenants.tsx
@@ -1,10 +1,11 @@
 import type { TenantInfo } from '@logto/schemas';
-import { useCallback, useEffect } from 'react';
+import { useCallback, useContext, useEffect } from 'react';
 
 import { useCloudApi } from '@/cloud/hooks/use-cloud-api';
 import { AppLoadingOffline } from '@/components/AppLoading/Offline';
 import Button from '@/components/Button';
 import DangerousRaw from '@/components/DangerousRaw';
+import { TenantsContext } from '@/contexts/TenantsProvider';
 
 import * as styles from './index.module.scss';
 
@@ -15,6 +16,7 @@ type Props = {
 
 const Tenants = ({ data, onAdd }: Props) => {
   const api = useCloudApi();
+  const { navigate } = useContext(TenantsContext);
 
   const createTenant = useCallback(async () => {
     onAdd(await api.post('api/tenants').json<TenantInfo>());
@@ -26,18 +28,25 @@ const Tenants = ({ data, onAdd }: Props) => {
     }
 
     if (data[0]) {
-      window.location.assign('/' + data[0].id);
+      navigate(data[0].id);
     } else {
       void createTenant();
     }
-  }, [createTenant, data]);
+  }, [createTenant, data, navigate]);
 
   if (data.length > 1) {
     return (
       <div className={styles.wrapper}>
         <h3>Choose a tenant</h3>
         {data.map(({ id }) => (
-          <a key={id} href={'/' + id}>
+          <a
+            key={id}
+            href={'/' + id}
+            onClick={(event) => {
+              event.preventDefault();
+              navigate(id);
+            }}
+          >
             <Button title={<DangerousRaw>{id}</DangerousRaw>} />
           </a>
         ))}
diff --git a/packages/console/src/cloud/pages/Main/index.tsx b/packages/console/src/cloud/pages/Main/index.tsx
index 1427bd1e5..4a1689399 100644
--- a/packages/console/src/cloud/pages/Main/index.tsx
+++ b/packages/console/src/cloud/pages/Main/index.tsx
@@ -5,7 +5,6 @@ import { useHref } from 'react-router-dom';
 
 import { useCloudApi } from '@/cloud/hooks/use-cloud-api';
 import { AppLoadingOffline } from '@/components/AppLoading/Offline';
-import { getUserTenantId } from '@/consts/tenants';
 import { TenantsContext } from '@/contexts/TenantsProvider';
 
 import Redirect from './Redirect';
@@ -13,7 +12,7 @@ import Tenants from './Tenants';
 
 const Protected = () => {
   const api = useCloudApi();
-  const { tenants, setTenants } = useContext(TenantsContext);
+  const { tenants, setTenants, currentTenantId } = useContext(TenantsContext);
 
   useEffect(() => {
     const loadTenants = async () => {
@@ -27,8 +26,6 @@ const Protected = () => {
   }, [api, setTenants, tenants]);
 
   if (tenants) {
-    const currentTenantId = getUserTenantId();
-
     if (currentTenantId) {
       return <Redirect tenants={tenants} toTenantId={currentTenantId} />;
     }
@@ -48,7 +45,8 @@ const Protected = () => {
 
 const Main = () => {
   const { isAuthenticated, isLoading, signIn } = useLogto();
-  const href = useHref(getUserTenantId() + '/callback');
+  const { currentTenantId } = useContext(TenantsContext);
+  const href = useHref(currentTenantId + '/callback');
 
   useEffect(() => {
     if (!isLoading && !isAuthenticated) {
diff --git a/packages/console/src/components/AppLoading/Offline.tsx b/packages/console/src/components/AppLoading/Offline.tsx
index 4e62602a6..88a13fdea 100644
--- a/packages/console/src/components/AppLoading/Offline.tsx
+++ b/packages/console/src/components/AppLoading/Offline.tsx
@@ -10,6 +10,9 @@ import { getTheme } from '@/utils/theme';
 
 import * as styles from './index.module.scss';
 
+/**
+ * An fullscreen loading component fetches local stored theme without sending request.
+ */
 export const AppLoadingOffline = () => {
   const theme = getTheme(
     trySafe(() => z.nativeEnum(AppearanceMode).parse(localStorage.getItem(themeStorageKey))) ??
diff --git a/packages/console/src/contexts/AppEndpointsProvider.tsx b/packages/console/src/contexts/AppEndpointsProvider.tsx
index 704675be7..fa07835c5 100644
--- a/packages/console/src/contexts/AppEndpointsProvider.tsx
+++ b/packages/console/src/contexts/AppEndpointsProvider.tsx
@@ -1,8 +1,10 @@
 import ky from 'ky';
 import type { ReactNode } from 'react';
-import { useMemo, useEffect, createContext, useState } from 'react';
+import { useContext, useMemo, useEffect, createContext, useState } from 'react';
 
-import { adminTenantEndpoint, getUserTenantId } from '@/consts';
+import { adminTenantEndpoint } from '@/consts';
+
+import { TenantsContext } from './TenantsProvider';
 
 type Props = {
   children: ReactNode;
@@ -17,24 +19,23 @@ export const AppEndpointsContext = createContext<AppEndpoints>({});
 
 const AppEndpointsProvider = ({ children }: Props) => {
   const [endpoints, setEndpoints] = useState<AppEndpoints>({});
+  const { currentTenantId } = useContext(TenantsContext);
   const memorizedContext = useMemo(() => endpoints, [endpoints]);
 
   useEffect(() => {
     const getEndpoint = async () => {
-      const tenantId = getUserTenantId();
-
-      if (!tenantId) {
+      if (!currentTenantId) {
         return;
       }
 
       const { user } = await ky
-        .get(new URL(`api/.well-known/endpoints/${tenantId}`, adminTenantEndpoint))
+        .get(new URL(`api/.well-known/endpoints/${currentTenantId}`, adminTenantEndpoint))
         .json<{ user: string }>();
       setEndpoints({ userEndpoint: new URL(user) });
     };
 
     void getEndpoint();
-  }, []);
+  }, [currentTenantId]);
 
   return (
     <AppEndpointsContext.Provider value={memorizedContext}>{children}</AppEndpointsContext.Provider>
diff --git a/packages/console/src/contexts/TenantsProvider.tsx b/packages/console/src/contexts/TenantsProvider.tsx
index 677d14f4f..0a7f005cb 100644
--- a/packages/console/src/contexts/TenantsProvider.tsx
+++ b/packages/console/src/contexts/TenantsProvider.tsx
@@ -2,9 +2,10 @@ import type { TenantInfo } from '@logto/schemas';
 import { defaultManagementApi } from '@logto/schemas';
 import { conditional, noop } from '@silverhand/essentials';
 import type { ReactNode } from 'react';
-import { useMemo, createContext, useState } from 'react';
+import { useCallback, useMemo, createContext, useState } from 'react';
 
 import { isCloud } from '@/consts/cloud';
+import { getUserTenantId } from '@/consts/tenants';
 
 type Props = {
   children: ReactNode;
@@ -15,6 +16,8 @@ export type Tenants = {
   isSettle: boolean;
   setTenants: (tenants: TenantInfo[]) => void;
   setIsSettle: (isSettle: boolean) => void;
+  currentTenantId: string;
+  navigate: (tenantId: string) => void;
 };
 
 const { tenantId, indicator } = defaultManagementApi.resource;
@@ -25,14 +28,23 @@ export const TenantsContext = createContext<Tenants>({
   setTenants: noop,
   isSettle: false,
   setIsSettle: noop,
+  currentTenantId: '',
+  navigate: noop,
 });
 
 const TenantsProvider = ({ children }: Props) => {
   const [tenants, setTenants] = useState(initialTenants);
   const [isSettle, setIsSettle] = useState(false);
+  const [currentTenantId, setCurrentTenantId] = useState(getUserTenantId());
+
+  const navigate = useCallback((tenantId: string) => {
+    window.history.pushState({}, '', '/' + tenantId);
+    setCurrentTenantId(tenantId);
+  }, []);
+
   const memorizedContext = useMemo(
-    () => ({ tenants, setTenants, isSettle, setIsSettle }),
-    [isSettle, tenants]
+    () => ({ tenants, setTenants, isSettle, setIsSettle, currentTenantId, navigate }),
+    [currentTenantId, isSettle, navigate, tenants]
   );
 
   return <TenantsContext.Provider value={memorizedContext}>{children}</TenantsContext.Provider>;
diff --git a/packages/console/src/hooks/use-api.ts b/packages/console/src/hooks/use-api.ts
index 186981ad9..a08bf8e24 100644
--- a/packages/console/src/hooks/use-api.ts
+++ b/packages/console/src/hooks/use-api.ts
@@ -5,8 +5,9 @@ import { useCallback, useContext, useMemo } from 'react';
 import { toast } from 'react-hot-toast';
 import { useTranslation } from 'react-i18next';
 
-import { getBasename, getManagementApi, getUserTenantId, requestTimeout } from '@/consts';
+import { getBasename, getManagementApi, requestTimeout } from '@/consts';
 import { AppEndpointsContext } from '@/contexts/AppEndpointsProvider';
+import { TenantsContext } from '@/contexts/TenantsProvider';
 
 import { useConfirmModal } from './use-confirm-modal';
 
@@ -30,11 +31,16 @@ type StaticApiProps = {
 export const useStaticApi = ({
   prefixUrl,
   hideErrorToast,
-  resourceIndicator = getManagementApi(getUserTenantId()).indicator,
+  resourceIndicator: resourceInput,
 }: StaticApiProps) => {
   const { isAuthenticated, getAccessToken, signOut } = useLogto();
+  const { currentTenantId } = useContext(TenantsContext);
   const { t, i18n } = useTranslation(undefined, { keyPrefix: 'admin_console' });
   const { show } = useConfirmModal();
+  const resourceIndicator = useMemo(
+    () => resourceInput ?? getManagementApi(currentTenantId).indicator,
+    [currentTenantId, resourceInput]
+  );
 
   const toastError = useCallback(
     async (response: Response) => {
diff --git a/packages/core/src/middleware/koa-body-etag.ts b/packages/core/src/middleware/koa-body-etag.ts
new file mode 100644
index 000000000..32cafbd22
--- /dev/null
+++ b/packages/core/src/middleware/koa-body-etag.ts
@@ -0,0 +1,29 @@
+import type { Nullable } from '@silverhand/essentials';
+import etag from 'etag';
+import type { Method } from 'got';
+import type { MiddlewareType } from 'koa';
+
+/**
+ * Create a middleware function that calculates the ETag value based on response body,
+ * and set status code to `304` and body to `null` if needed.
+ *
+ * @param methods An array of methods to match
+ * @default methods ['GET']
+ */
+export default function koaBodyEtag<StateT, ContextT, ResponseBodyT>(
+  methods: Array<Uppercase<Method>> = ['GET']
+): MiddlewareType<StateT, ContextT, Nullable<ResponseBodyT>> {
+  return async (ctx, next) => {
+    await next();
+
+    // eslint-disable-next-line no-restricted-syntax
+    if (methods.includes(ctx.method as Uppercase<Method>)) {
+      ctx.response.etag = etag(JSON.stringify(ctx.body));
+
+      if (ctx.fresh) {
+        ctx.status = 304;
+        ctx.body = null;
+      }
+    }
+  };
+}
diff --git a/packages/core/src/oidc/init.ts b/packages/core/src/oidc/init.ts
index 6f718eb18..37bfdbeb0 100644
--- a/packages/core/src/oidc/init.ts
+++ b/packages/core/src/oidc/init.ts
@@ -12,6 +12,7 @@ import snakecaseKeys from 'snakecase-keys';
 import type { EnvSet } from '#src/env-set/index.js';
 import { addOidcEventListeners } from '#src/event-listeners/index.js';
 import koaAuditLog from '#src/middleware/koa-audit-log.js';
+import koaBodyEtag from '#src/middleware/koa-body-etag.js';
 import postgresAdapter from '#src/oidc/adapter.js';
 import { isOriginAllowed, validateCustomClientMetadata } from '#src/oidc/utils.js';
 import { routes } from '#src/routes/consts.js';
@@ -234,6 +235,7 @@ export default function initOidc(envSet: EnvSet, queries: Queries, libraries: Li
 
   // Provide audit log context for event listeners
   oidc.use(koaAuditLog(queries));
+  oidc.use(koaBodyEtag());
 
   return oidc;
 }
diff --git a/packages/core/src/routes/well-known.ts b/packages/core/src/routes/well-known.ts
index 352dfcd13..d00efb0f2 100644
--- a/packages/core/src/routes/well-known.ts
+++ b/packages/core/src/routes/well-known.ts
@@ -1,10 +1,10 @@
 import type { ConnectorMetadata } from '@logto/connector-kit';
 import { ConnectorType } from '@logto/connector-kit';
 import { adminTenantId } from '@logto/schemas';
-import etag from 'etag';
 
 import { EnvSet, getTenantEndpoint } from '#src/env-set/index.js';
 import RequestError from '#src/errors/RequestError/index.js';
+import koaBodyEtag from '#src/middleware/koa-body-etag.js';
 
 import type { AnonymousRouter, RouterInitArgs } from './types.js';
 
@@ -30,49 +30,38 @@ export default function wellKnownRoutes<T extends AnonymousRouter>(
     });
   }
 
-  router.get(
-    '/.well-known/sign-in-exp',
-    async (ctx, next) => {
-      const [signInExperience, logtoConnectors] = await Promise.all([
-        getSignInExperience(),
-        getLogtoConnectors(),
-      ]);
+  router.get('/.well-known/sign-in-exp', async (ctx, next) => {
+    const [signInExperience, logtoConnectors] = await Promise.all([
+      getSignInExperience(),
+      getLogtoConnectors(),
+    ]);
 
-      const forgotPassword = {
-        phone: logtoConnectors.some(({ type }) => type === ConnectorType.Sms),
-        email: logtoConnectors.some(({ type }) => type === ConnectorType.Email),
-      };
+    const forgotPassword = {
+      phone: logtoConnectors.some(({ type }) => type === ConnectorType.Sms),
+      email: logtoConnectors.some(({ type }) => type === ConnectorType.Email),
+    };
 
-      const socialConnectors = signInExperience.socialSignInConnectorTargets.reduce<
-        Array<ConnectorMetadata & { id: string }>
-      >((previous, connectorTarget) => {
-        const connectors = logtoConnectors.filter(
-          ({ metadata: { target } }) => target === connectorTarget
-        );
+    const socialConnectors = signInExperience.socialSignInConnectorTargets.reduce<
+      Array<ConnectorMetadata & { id: string }>
+    >((previous, connectorTarget) => {
+      const connectors = logtoConnectors.filter(
+        ({ metadata: { target } }) => target === connectorTarget
+      );
 
-        return [
-          ...previous,
-          ...connectors.map(({ metadata, dbEntry: { id } }) => ({ ...metadata, id })),
-        ];
-      }, []);
+      return [
+        ...previous,
+        ...connectors.map(({ metadata, dbEntry: { id } }) => ({ ...metadata, id })),
+      ];
+    }, []);
 
-      ctx.body = {
-        ...signInExperience,
-        socialConnectors,
-        forgotPassword,
-      };
+    ctx.body = {
+      ...signInExperience,
+      socialConnectors,
+      forgotPassword,
+    };
 
-      return next();
-    },
-    async (ctx, next) => {
-      await next();
+    return next();
+  });
 
-      ctx.response.etag = etag(JSON.stringify(ctx.body));
-
-      if (ctx.fresh) {
-        ctx.status = 304;
-        ctx.body = null;
-      }
-    }
-  );
+  router.use(koaBodyEtag());
 }