From 6dd90ce9cab1099f73f1b60ca33a746f3a6108cf Mon Sep 17 00:00:00 2001
From: Mike Cao <mike@mikecao.com>
Date: Tue, 4 Apr 2023 23:29:54 -0700
Subject: [PATCH] Refactored SSO process.

---
 lib/auth.ts             | 21 ++++++++++++++++++++-
 pages/api/auth/login.ts | 10 +++-------
 pages/api/auth/sso.ts   | 18 ++++++++++++++++++
 pages/sso.js            | 16 ++++++++++++++--
 4 files changed, 55 insertions(+), 10 deletions(-)
 create mode 100644 pages/api/auth/sso.ts

diff --git a/lib/auth.ts b/lib/auth.ts
index 6cc48ff2..7c83034c 100644
--- a/lib/auth.ts
+++ b/lib/auth.ts
@@ -1,8 +1,15 @@
 import debug from 'debug';
+import redis from '@umami/redis-client';
 import cache from 'lib/cache';
 import { PERMISSIONS, ROLE_PERMISSIONS, SHARE_TOKEN_HEADER } from 'lib/constants';
 import { secret } from 'lib/crypto';
-import { ensureArray, parseSecureToken, parseToken } from 'next-basics';
+import {
+  createSecureToken,
+  ensureArray,
+  getRandomChars,
+  parseSecureToken,
+  parseToken,
+} from 'next-basics';
 import { getTeamUser, getTeamUserById } from 'queries';
 import { getTeamWebsite, getTeamWebsiteByTeamMemberId } from 'queries/admin/teamWebsite';
 import { validate } from 'uuid';
@@ -11,6 +18,18 @@ import { loadWebsite } from './query';
 
 const log = debug('umami:auth');
 
+export async function setAuthKey(user, expire = 0) {
+  const authKey = `auth:${getRandomChars(32)}`;
+
+  await redis.set(authKey, user);
+
+  if (expire) {
+    await redis.expire(authKey, expire);
+  }
+
+  return createSecureToken({ authKey }, secret());
+}
+
 export function getAuthToken(req) {
   try {
     return req.headers.authorization.split(' ')[1];
diff --git a/pages/api/auth/login.ts b/pages/api/auth/login.ts
index 0e226346..81e2df5b 100644
--- a/pages/api/auth/login.ts
+++ b/pages/api/auth/login.ts
@@ -1,3 +1,4 @@
+import { NextApiResponse } from 'next';
 import {
   ok,
   unauthorized,
@@ -5,13 +6,12 @@ import {
   checkPassword,
   createSecureToken,
   methodNotAllowed,
-  getRandomChars,
 } from 'next-basics';
 import redis from '@umami/redis-client';
 import { getUser } from 'queries';
 import { secret } from 'lib/crypto';
 import { NextApiRequestQueryBody, User } from 'lib/types';
-import { NextApiResponse } from 'next';
+import { setAuthKey } from 'lib/auth';
 
 export interface LoginRequestBody {
   username: string;
@@ -38,11 +38,7 @@ export default async (
 
     if (user && checkPassword(password, user.password)) {
       if (redis.enabled) {
-        const authKey = `auth:${getRandomChars(32)}`;
-
-        await redis.set(authKey, user);
-
-        const token = createSecureToken({ authKey }, secret());
+        const token = await setAuthKey(user);
 
         return ok(res, { token, user });
       }
diff --git a/pages/api/auth/sso.ts b/pages/api/auth/sso.ts
new file mode 100644
index 00000000..66cbd8e2
--- /dev/null
+++ b/pages/api/auth/sso.ts
@@ -0,0 +1,18 @@
+import { NextApiRequestAuth } from 'lib/types';
+import { useAuth } from 'lib/middleware';
+import { NextApiResponse } from 'next';
+import { badRequest, ok } from 'next-basics';
+import redis from '@umami/redis-client';
+import { setAuthKey } from 'lib/auth';
+
+export default async (req: NextApiRequestAuth, res: NextApiResponse) => {
+  await useAuth(req, res);
+
+  if (redis.enabled && req.auth.user) {
+    const token = await setAuthKey(req.auth.user, 86400);
+
+    return ok(res, { user: req.auth.user, token });
+  }
+
+  return badRequest(res);
+};
diff --git a/pages/sso.js b/pages/sso.js
index c3d499c3..bc978d2c 100644
--- a/pages/sso.js
+++ b/pages/sso.js
@@ -2,18 +2,30 @@ import { useEffect } from 'react';
 import { Loading } from 'react-basics';
 import { useRouter } from 'next/router';
 import { setClientAuthToken } from 'lib/client';
+import useApi from '../hooks/useApi';
 
 export default function SingleSignOnPage() {
+  const { post } = useApi();
   const router = useRouter();
   const { token, url } = router.query;
 
   useEffect(() => {
+    const signOn = async token => {
+      setClientAuthToken(token);
+
+      const data = await post('/auth/sso');
+
+      setClientAuthToken(data.token);
+
+      await router.push(url);
+    };
+
     if (url && token) {
       setClientAuthToken(token);
 
-      router.push(url);
+      signOn();
     }
-  }, [router, url, token]);
+  }, [url, token]);
 
   return <Loading size="xl" />;
 }