From 498daef9d30c922b6e355574e76003bd39bb1099 Mon Sep 17 00:00:00 2001
From: Keith Stevens <fozziethebeat@gmail.com>
Date: Fri, 6 Jan 2023 20:01:02 +0900
Subject: [PATCH] Creating a role field for Users and seeding it everytime a
 user logs in based on a shortlist

---
 website/.env                                |  2 +
 website/prisma/schema.prisma                |  1 +
 website/src/pages/api/auth/[...nextauth].ts | 66 +++++++++++++++++++++
 3 files changed, 69 insertions(+)

diff --git a/website/.env b/website/.env
index 9544836b..65d8b88e 100644
--- a/website/.env
+++ b/website/.env
@@ -1,3 +1,5 @@
+ADMIN_USERS = "credentials:admin,discord:root,email:admin@example.com"
+
 # The database created by running the jobs in /scripts/frontend-development/docker-compose.yaml
 DATABASE_URL=postgres://postgres:postgres@localhost:5433/oasst_web
 
diff --git a/website/prisma/schema.prisma b/website/prisma/schema.prisma
index 3e379d43..f9eab3b7 100644
--- a/website/prisma/schema.prisma
+++ b/website/prisma/schema.prisma
@@ -41,6 +41,7 @@ model User {
   email         String?   @unique
   emailVerified DateTime?
   image         String?
+  role          String    @default("general")
 
   accounts      Account[]
   sessions      Session[]
diff --git a/website/src/pages/api/auth/[...nextauth].ts b/website/src/pages/api/auth/[...nextauth].ts
index 8614de97..f69be040 100644
--- a/website/src/pages/api/auth/[...nextauth].ts
+++ b/website/src/pages/api/auth/[...nextauth].ts
@@ -59,6 +59,17 @@ if (boolean(process.env.DEBUG_LOGIN) || process.env.NODE_ENV === "development")
   );
 }
 
+// Create a map of provider types to a set of admin user identifiers based on
+// the environment variables.  We assume the list is separated by ',' and each
+// entry is separated by ':'.
+const adminUserMap = process.env.ADMIN_USERS.split(",").reduce((result, entry) => {
+  const [authType, id] = entry.split(":");
+  const s = result.get(authType) || new Set();
+  s.add(id);
+  result.set(authType, s);
+  return result;
+}, new Map());
+
 export const authOptions: AuthOptions = {
   // Ensure we can store user data in a database.
   adapter: PrismaAdapter(prisma),
@@ -68,6 +79,61 @@ export const authOptions: AuthOptions = {
     verifyRequest: "/auth/verify",
     // error: "/auth/error", -Will be used later
   },
+  callbacks: {
+    /**
+     * Ensure we propagate the user's role when creating the session from the
+     * token.
+     */
+    async session({ session, user, token }) {
+      session.user.role = token.role;
+      return session;
+    },
+    /**
+     * When creating a token, fetch the user's role and inject it in the token.
+     * This let's use forward the role to the session object.
+     */
+    async jwt({ token, profile, account }) {
+      const { role } = await prisma.user.findUnique({
+        where: { id: token.sub },
+        select: { role: true },
+      });
+      token.role = role;
+      return token;
+    },
+  },
+  events: {
+    /**
+     * Update the user's role after they have successfully signed in
+     */
+    async signIn({ user, account }) {
+      console.log(account.provider);
+      // Get the admin list for the user's auth type.
+      const adminForAccountType = adminUserMap.get(account.provider);
+
+      // Return early if there's no admin list.
+      if (!adminForAccountType) {
+        return;
+      }
+
+      // Return early if the user is already an admin to reduce database
+      // writes.
+      if (user?.role === "admin") {
+        return;
+      }
+
+      // Update the database if the user is an admin.
+      if (adminForAccountType.has(account.providerAccountId)) {
+        await prisma.user.update({
+          data: {
+            role: "admin",
+          },
+          where: {
+            id: user.id,
+          },
+        });
+      }
+    },
+  },
   session: {
     strategy: "jwt",
   },