From dc71a77d8cf5aeaf509738002b08bfb9c0ab3880 Mon Sep 17 00:00:00 2001
From: YANG QIA <2013xile@gmail.com>
Date: Sun, 14 Jan 2024 22:26:13 +0800
Subject: [PATCH] fix(auth): redirect URL after signing in by SSO sucessfully
 (#3387)

* fix: fix T-2900

* fix: saml
---
 .../@nocobase/plugin-cas/src/client/SigninPage.tsx        | 8 ++++----
 .../plugins/@nocobase/plugin-cas/src/client/index.tsx     | 2 +-
 .../@nocobase/plugin-cas/src/server/actions/login.ts      | 6 ++++--
 .../@nocobase/plugin-cas/src/server/actions/service.ts    | 6 +++---
 .../@nocobase/plugin-oidc/src/client/OIDCButton.tsx       | 6 +++++-
 .../plugin-oidc/src/server/actions/getAuthUrl.ts          | 5 +++--
 .../@nocobase/plugin-oidc/src/server/actions/redirect.ts  | 5 +++--
 .../@nocobase/plugin-saml/src/client/SAMLButton.tsx       | 6 +++++-
 .../plugin-saml/src/server/actions/getAuthUrl.ts          | 3 ++-
 .../@nocobase/plugin-saml/src/server/actions/redirect.ts  | 5 +++--
 .../plugins/@nocobase/plugin-saml/src/server/saml-auth.ts | 6 +++---
 11 files changed, 36 insertions(+), 22 deletions(-)

diff --git a/packages/plugins/@nocobase/plugin-cas/src/client/SigninPage.tsx b/packages/plugins/@nocobase/plugin-cas/src/client/SigninPage.tsx
index 1956ac699..0918d46b0 100644
--- a/packages/plugins/@nocobase/plugin-cas/src/client/SigninPage.tsx
+++ b/packages/plugins/@nocobase/plugin-cas/src/client/SigninPage.tsx
@@ -6,17 +6,17 @@ import { getSubAppName } from '@nocobase/sdk';
 import { Authenticator } from '@nocobase/plugin-auth/client';
 
 export const SigninPage = (props: { authenticator: Authenticator }) => {
-  const location = useLocation();
-
   const authenticator = props.authenticator;
+  const location = useLocation();
+  const params = new URLSearchParams(location.search);
+  const redirect = params.get('redirect');
 
   const app = getSubAppName() || 'main';
   const login = async () => {
-    window.location.replace(`/api/cas:login?authenticator=${authenticator.name}&__appName=${app}`);
+    window.location.replace(`/api/cas:login?authenticator=${authenticator.name}&__appName=${app}&redirect=${redirect}`);
   };
 
   useEffect(() => {
-    const params = new URLSearchParams(location.search);
     const name = params.get('authenticator');
     const error = params.get('error');
     if (name !== authenticator.name) {
diff --git a/packages/plugins/@nocobase/plugin-cas/src/client/index.tsx b/packages/plugins/@nocobase/plugin-cas/src/client/index.tsx
index 51bb1df36..bf5cb1e8e 100644
--- a/packages/plugins/@nocobase/plugin-cas/src/client/index.tsx
+++ b/packages/plugins/@nocobase/plugin-cas/src/client/index.tsx
@@ -9,7 +9,7 @@ export class SamlPlugin extends Plugin {
     const auth = this.app.pm.get(AuthPlugin);
     auth.registerType(authType, {
       components: {
-        SignInForm: SigninPage,
+        SignInButton: SigninPage,
         AdminSettingsForm: Options,
       },
     });
diff --git a/packages/plugins/@nocobase/plugin-cas/src/server/actions/login.ts b/packages/plugins/@nocobase/plugin-cas/src/server/actions/login.ts
index f563da933..51ab5e99c 100644
--- a/packages/plugins/@nocobase/plugin-cas/src/server/actions/login.ts
+++ b/packages/plugins/@nocobase/plugin-cas/src/server/actions/login.ts
@@ -2,10 +2,12 @@ import { Context, Next } from '@nocobase/actions';
 import { CASAuth } from '../auth';
 
 export const login = async (ctx: Context, next: Next) => {
-  const { authenticator } = ctx.action.params;
+  const { authenticator, redirect = '' } = ctx.action.params;
   const auth = (await ctx.app.authManager.get(authenticator, ctx)) as CASAuth;
   const { casUrl, serviceUrl } = auth.getOptions();
-  const service = encodeURIComponent(`${serviceUrl}?authenticator=${authenticator}&__appName=${ctx.app.name}`);
+  const service = encodeURIComponent(
+    `${serviceUrl}?authenticator=${authenticator}&__appName=${ctx.app.name}&redirect=${redirect}`,
+  );
   ctx.redirect(`${casUrl}/login?service=${service}`);
   next();
 };
diff --git a/packages/plugins/@nocobase/plugin-cas/src/server/actions/service.ts b/packages/plugins/@nocobase/plugin-cas/src/server/actions/service.ts
index e42fdc19a..132faf50c 100644
--- a/packages/plugins/@nocobase/plugin-cas/src/server/actions/service.ts
+++ b/packages/plugins/@nocobase/plugin-cas/src/server/actions/service.ts
@@ -3,7 +3,7 @@ import { AppSupervisor } from '@nocobase/server';
 import { CASAuth } from '../auth';
 
 export const service = async (ctx: Context, next: Next) => {
-  const { authenticator, __appName: appName } = ctx.action.params;
+  const { authenticator, __appName: appName, redirect = '/admin' } = ctx.action.params;
 
   let prefix = '';
   if (appName && appName !== 'main') {
@@ -16,9 +16,9 @@ export const service = async (ctx: Context, next: Next) => {
   const auth = (await ctx.app.authManager.get(authenticator, ctx)) as CASAuth;
   try {
     const { token } = await auth.signIn();
-    ctx.redirect(`${prefix}/admin?authenticator=${authenticator}&token=${token}`);
+    ctx.redirect(`${prefix}${redirect}?authenticator=${authenticator}&token=${token}`);
   } catch (error) {
-    ctx.redirect(`${prefix}/signin?authenticator=${authenticator}&error=${error.message}`);
+    ctx.redirect(`${prefix}/signin?authenticator=${authenticator}&error=${error.message}&redirect=${redirect}`);
   }
   return next();
 };
diff --git a/packages/plugins/@nocobase/plugin-oidc/src/client/OIDCButton.tsx b/packages/plugins/@nocobase/plugin-oidc/src/client/OIDCButton.tsx
index 9e393ffc7..1246d51ac 100644
--- a/packages/plugins/@nocobase/plugin-oidc/src/client/OIDCButton.tsx
+++ b/packages/plugins/@nocobase/plugin-oidc/src/client/OIDCButton.tsx
@@ -15,6 +15,8 @@ export const OIDCButton = ({ authenticator }: { authenticator: Authenticator })
   const { t } = useOidcTranslation();
   const api = useAPIClient();
   const location = useLocation();
+  const params = new URLSearchParams(location.search);
+  const redirect = params.get('redirect');
 
   const login = async () => {
     const response = await api.request({
@@ -23,6 +25,9 @@ export const OIDCButton = ({ authenticator }: { authenticator: Authenticator })
       headers: {
         'X-Authenticator': authenticator.name,
       },
+      data: {
+        redirect,
+      },
     });
 
     const authUrl = response?.data?.data;
@@ -30,7 +35,6 @@ export const OIDCButton = ({ authenticator }: { authenticator: Authenticator })
   };
 
   useEffect(() => {
-    const params = new URLSearchParams(location.search);
     const name = params.get('authenticator');
     const error = params.get('error');
     if (name !== authenticator.name) {
diff --git a/packages/plugins/@nocobase/plugin-oidc/src/server/actions/getAuthUrl.ts b/packages/plugins/@nocobase/plugin-oidc/src/server/actions/getAuthUrl.ts
index 67901c204..c96f6e25c 100644
--- a/packages/plugins/@nocobase/plugin-oidc/src/server/actions/getAuthUrl.ts
+++ b/packages/plugins/@nocobase/plugin-oidc/src/server/actions/getAuthUrl.ts
@@ -4,6 +4,7 @@ import { nanoid } from 'nanoid';
 import { cookieName } from '../../constants';
 
 export const getAuthUrl = async (ctx: Context, next: Next) => {
+  const { redirect = '' } = ctx.action.params.values || {};
   const app = ctx.app.name;
   const auth = ctx.auth as OIDCAuth;
   const client = await auth.createOIDCClient();
@@ -16,8 +17,8 @@ export const getAuthUrl = async (ctx: Context, next: Next) => {
   ctx.body = client.authorizationUrl({
     response_type: 'code',
     scope: scope || 'openid email profile',
-    redirect_uri: auth.getRedirectUri(),
-    state: encodeURIComponent(`token=${token}&name=${ctx.headers['x-authenticator']}&app=${app}`),
+    redirect_uri: `${auth.getRedirectUri()}`,
+    state: encodeURIComponent(`token=${token}&name=${ctx.headers['x-authenticator']}&app=${app}&redirect=${redirect}`),
   });
 
   return next();
diff --git a/packages/plugins/@nocobase/plugin-oidc/src/server/actions/redirect.ts b/packages/plugins/@nocobase/plugin-oidc/src/server/actions/redirect.ts
index fb642bd65..42d6b04f0 100644
--- a/packages/plugins/@nocobase/plugin-oidc/src/server/actions/redirect.ts
+++ b/packages/plugins/@nocobase/plugin-oidc/src/server/actions/redirect.ts
@@ -9,6 +9,7 @@ export const redirect = async (ctx: Context, next: Next) => {
   const search = new URLSearchParams(decodeURIComponent(state));
   const authenticator = search.get('name');
   const appName = search.get('app');
+  const redirect = search.get('redirect') || '/admin';
   let prefix = '';
   if (appName && appName !== 'main') {
     const appSupervisor = AppSupervisor.getInstance();
@@ -19,10 +20,10 @@ export const redirect = async (ctx: Context, next: Next) => {
   const auth = (await ctx.app.authManager.get(authenticator, ctx)) as OIDCAuth;
   try {
     const { token } = await auth.signIn();
-    ctx.redirect(`${prefix}/admin?authenticator=${authenticator}&token=${token}`);
+    ctx.redirect(`${prefix}${redirect}?authenticator=${authenticator}&token=${token}`);
   } catch (error) {
     ctx.logger.error('OIDC auth error', { error });
-    ctx.redirect(`${prefix}/signin?authenticator=${authenticator}&error=${error.message}`);
+    ctx.redirect(`${prefix}/signin?redirect=${redirect}&authenticator=${authenticator}&error=${error.message}`);
   }
   await next();
 };
diff --git a/packages/plugins/@nocobase/plugin-saml/src/client/SAMLButton.tsx b/packages/plugins/@nocobase/plugin-saml/src/client/SAMLButton.tsx
index 21653c09c..ad88a2e6a 100644
--- a/packages/plugins/@nocobase/plugin-saml/src/client/SAMLButton.tsx
+++ b/packages/plugins/@nocobase/plugin-saml/src/client/SAMLButton.tsx
@@ -10,6 +10,8 @@ export const SAMLButton = ({ authenticator }: { authenticator: Authenticator })
   const { t } = useSamlTranslation();
   const api = useAPIClient();
   const location = useLocation();
+  const params = new URLSearchParams(location.search);
+  const redirect = params.get('redirect');
 
   const login = async () => {
     const response = await api.request({
@@ -18,6 +20,9 @@ export const SAMLButton = ({ authenticator }: { authenticator: Authenticator })
       headers: {
         'X-Authenticator': authenticator.name,
       },
+      data: {
+        redirect,
+      },
     });
 
     const authUrl = response?.data?.data;
@@ -25,7 +30,6 @@ export const SAMLButton = ({ authenticator }: { authenticator: Authenticator })
   };
 
   useEffect(() => {
-    const params = new URLSearchParams(location.search);
     const name = params.get('authenticator');
     const error = params.get('error');
     if (name !== authenticator.name) {
diff --git a/packages/plugins/@nocobase/plugin-saml/src/server/actions/getAuthUrl.ts b/packages/plugins/@nocobase/plugin-saml/src/server/actions/getAuthUrl.ts
index 136f475eb..c3fad8eb3 100644
--- a/packages/plugins/@nocobase/plugin-saml/src/server/actions/getAuthUrl.ts
+++ b/packages/plugins/@nocobase/plugin-saml/src/server/actions/getAuthUrl.ts
@@ -3,11 +3,12 @@ import { SAML } from '@node-saml/node-saml';
 import { SAMLAuth } from '../saml-auth';
 
 export const getAuthUrl = async (ctx: Context, next: Next) => {
+  const { redirect = '' } = ctx.action.params.values || {};
   const auth = ctx.auth as SAMLAuth;
   const options = auth.getOptions();
   const saml = new SAML(options);
 
-  ctx.body = await saml.getAuthorizeUrlAsync('', '', {});
+  ctx.body = await saml.getAuthorizeUrlAsync(redirect, '', {});
 
   return next();
 };
diff --git a/packages/plugins/@nocobase/plugin-saml/src/server/actions/redirect.ts b/packages/plugins/@nocobase/plugin-saml/src/server/actions/redirect.ts
index 41e54144b..0acddaf53 100644
--- a/packages/plugins/@nocobase/plugin-saml/src/server/actions/redirect.ts
+++ b/packages/plugins/@nocobase/plugin-saml/src/server/actions/redirect.ts
@@ -4,6 +4,7 @@ import { AppSupervisor } from '@nocobase/server';
 
 export const redirect = async (ctx: Context, next: Next) => {
   const { authenticator, __appName: appName } = ctx.action.params || {};
+  const { RelayState: redirect = '/admin' } = ctx.action.params.values || {};
   let prefix = '';
   if (appName && appName !== 'main') {
     const appSupervisor = AppSupervisor.getInstance();
@@ -14,9 +15,9 @@ export const redirect = async (ctx: Context, next: Next) => {
   const auth = (await ctx.app.authManager.get(authenticator, ctx)) as SAMLAuth;
   try {
     const { token } = await auth.signIn();
-    ctx.redirect(`${prefix}/admin?authenticator=${authenticator}&token=${token}`);
+    ctx.redirect(`${prefix}${redirect}?authenticator=${authenticator}&token=${token}`);
   } catch (error) {
-    ctx.redirect(`${prefix}/signin?authenticator=${authenticator}&error=${error.message}`);
+    ctx.redirect(`${prefix}/signin?authenticator=${authenticator}&error=${error.message}&redirect=${redirect}`);
   }
   await next();
 };
diff --git a/packages/plugins/@nocobase/plugin-saml/src/server/saml-auth.ts b/packages/plugins/@nocobase/plugin-saml/src/server/saml-auth.ts
index ae33a0cde..10ffee08c 100644
--- a/packages/plugins/@nocobase/plugin-saml/src/server/saml-auth.ts
+++ b/packages/plugins/@nocobase/plugin-saml/src/server/saml-auth.ts
@@ -18,19 +18,19 @@ export class SAMLAuth extends BaseAuth {
     });
   }
 
-  getOptions() {
+  getOptions(): SamlConfig {
     const ctx = this.ctx;
     const { ssoUrl, certificate, idpIssuer, http }: SAMLOptions = this.options?.saml || {};
     const name = this.authenticator.get('name');
     const protocol = http ? 'http' : 'https';
     return {
-      callbackUrl: `${protocol}://${ctx.host}/api/saml:redirect?authenticator=${name}&app=${ctx.app.name}`,
+      callbackUrl: `${protocol}://${ctx.host}/api/saml:redirect?authenticator=${name}&__appName=${ctx.app.name}`,
       entryPoint: ssoUrl,
       issuer: name,
       cert: certificate,
       idpIssuer,
       wantAssertionsSigned: false,
-    } as SamlConfig;
+    };
   }
 
   async validate() {