feat: move to whitelist database table from growthbook (#864)

* feat: move to whitelist database table from growthbook

* fix: use lowercase email in tests

* chore: add additional guard clause for email validation

* fix: enforce using last part of email address

apps/studio/prisma/generated/generatedTypes.ts

@@ -114,6 +114,13 @@ export interface Version {
   publishedBy: string
   updatedAt: Generated<Timestamp>
 }
+export interface Whitelist {
+  id: GeneratedAlways<number>
+  email: string
+  expiry: Timestamp | null
+  createdAt: Generated<Timestamp>
+  updatedAt: Generated<Timestamp>
+}
 export interface DB {
   Blob: Blob
   Footer: Footer
@@ -126,4 +133,5 @@ export interface DB {
   User: User
   VerificationToken: VerificationToken
   Version: Version
+  Whitelist: Whitelist
 }

apps/studio/prisma/generated/selectableTypes.ts

@@ -14,3 +14,4 @@ export type SiteMember = Selectable<T.SiteMember>
 export type User = Selectable<T.User>
 export type VerificationToken = Selectable<T.VerificationToken>
 export type Version = Selectable<T.Version>
+export type Whitelist = Selectable<T.Whitelist>

apps/studio/prisma/migrations/20241107033200_add_whitelist_table/migration.sql

@@ -0,0 +1,19 @@
+-- CreateTable
+CREATE TABLE "Whitelist" (
+    "id" SERIAL NOT NULL,
+    "email" TEXT NOT NULL,
+    "expiry" TIMESTAMP(3),
+    "createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updatedAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+
+    CONSTRAINT "Whitelist_pkey" PRIMARY KEY ("id")
+);
+
+-- CreateIndex
+CREATE UNIQUE INDEX "Whitelist_email_key" ON "Whitelist"("email");
+
+-- CreateIndex
+CREATE INDEX "Whitelist_email_idx" ON "Whitelist"("email");
+
+-- AlterTable
+CREATE TRIGGER update_timestamp BEFORE UPDATE ON "Whitelist" FOR EACH ROW EXECUTE PROCEDURE moddatetime("updatedAt");

apps/studio/prisma/schema.prisma

@@ -212,3 +212,13 @@ model RateLimiterFlexible {
   points Int
   expire DateTime?
 }
+
+model Whitelist {
+  id        Int       @id @default(autoincrement())
+  email     String    @unique
+  expiry    DateTime?
+  createdAt DateTime  @default(now())
+  updatedAt DateTime  @default(now()) @updatedAt
+
+  @@index([email])
+}

apps/studio/prisma/seed.ts

@@ -259,6 +259,18 @@ async function main() {
     )
     .executeTakeFirstOrThrow()
 
+  await db
+    .insertInto("Whitelist")
+    .values({
+      email: "@open.gov.sg",
+    })
+    .onConflict((oc) =>
+      oc
+        .column("email")
+        .doUpdateSet((eb) => ({ email: eb.ref("excluded.email") })),
+    )
+    .executeTakeFirstOrThrow()
+
   await Promise.all(
     [...ISOMER_ADMINS, ...ISOMER_MIGRATORS, EDITOR_USER, PUBLISHER_USER].map(
       async (name) => {

apps/studio/src/server/modules/auth/email/__tests__/email.router.test.ts

@@ -3,6 +3,7 @@ import {
   applySession,
   createMockRequest,
 } from "tests/integration/helpers/iron-session"
+import { setUpWhitelist } from "tests/integration/helpers/seed"
 import { describe, expect, it } from "vitest"
 
 import { env } from "~/env.mjs"
@@ -15,16 +16,17 @@ import { getIpFingerprint, LOCALHOST } from "../utils"
 describe("auth.email", () => {
   let caller: Awaited<ReturnType<typeof emailSessionRouter.createCaller>>
   let session: ReturnType<typeof applySession>
+  const TEST_VALID_EMAIL = "[email protected]"
 
   beforeEach(async () => {
-    await resetTables("User", "VerificationToken")
+    await resetTables("User", "VerificationToken", "Whitelist")
+    await setUpWhitelist({ email: TEST_VALID_EMAIL })
     session = applySession()
     const ctx = createMockRequest(session)
     caller = emailSessionRouter.createCaller(ctx)
   })
 
   describe("login", () => {
-    const TEST_VALID_EMAIL = "[email protected]"
     it("should throw if email is not provided", async () => {
       // Act
       const result = caller.login({ email: "" })
@@ -72,7 +74,6 @@ describe("auth.email", () => {
   })
 
   describe("verifyOtp", () => {
-    const TEST_VALID_EMAIL = "[email protected]"
     const VALID_OTP = "123456"
     const VALID_TOKEN_HASH = createTokenHash(VALID_OTP, TEST_VALID_EMAIL)
     const INVALID_OTP = "987643"

apps/studio/src/server/modules/auth/email/email.router.ts

@@ -10,6 +10,7 @@ import {
 import { publicProcedure, router } from "~/server/trpc"
 import { getBaseUrl } from "~/utils/getBaseUrl"
 import { defaultMeSelect } from "../../me/me.select"
+import { isEmailWhitelisted } from "../../whitelist/whitelist.service"
 import { VerificationError } from "../auth.error"
 import { verifyToken } from "../auth.service"
 import { createTokenHash, createVfnPrefix, createVfnToken } from "../auth.util"
@@ -21,19 +22,13 @@ export const emailSessionRouter = router({
     .input(emailSignInSchema)
     .meta({ rateLimitOptions: {} })
     .mutation(async ({ ctx, input: { email } }) => {
-      if (env.NODE_ENV === "production") {
-        // check if whitelisted email on Growthbook
-        const defaultWhitelist: string[] = []
-        const whitelistedUsers = ctx.gb.getFeatureValue("whitelisted_users", {
-          whitelist: defaultWhitelist,
-        })
+      const isWhitelisted = await isEmailWhitelisted(email)
 
-        if (!whitelistedUsers.whitelist.includes(email)) {
-          throw new TRPCError({
-            code: "UNAUTHORIZED",
-            message: "Unauthorized. Contact Isomer support.",
-          })
-        }
+      if (!isWhitelisted) {
+        throw new TRPCError({
+          code: "UNAUTHORIZED",
+          message: "Unauthorized. Contact Isomer support.",
+        })
       }
 
       // TODO: instead of storing expires, store issuedAt to calculate when the next otp can be re-issued

apps/studio/src/server/modules/whitelist/__tests__/whitelist.service.test.ts

@@ -0,0 +1,125 @@
+import { resetTables } from "tests/integration/helpers/db"
+import { setUpWhitelist } from "tests/integration/helpers/seed"
+
+import { isEmailWhitelisted } from "../whitelist.service"
+
+describe("whitelist.service", () => {
+  beforeAll(async () => {
+    const oneYearFromNow = new Date()
+    oneYearFromNow.setFullYear(oneYearFromNow.getFullYear() + 1)
+    const oneYearAgo = new Date()
+    oneYearAgo.setFullYear(oneYearAgo.getFullYear() - 1)
+
+    await resetTables("Whitelist")
+    await setUpWhitelist({ email: "[email protected]" })
+    await setUpWhitelist({
+      email: "[email protected]",
+      expiry: oneYearFromNow,
+    })
+    await setUpWhitelist({
+      email: "[email protected]",
+      expiry: oneYearAgo,
+    })
+    await setUpWhitelist({ email: ".gov.sg" })
+    await setUpWhitelist({
+      email: "@vendor.com.sg",
+    })
+    await setUpWhitelist({
+      email: "@whitelisted.com.sg",
+      expiry: oneYearFromNow,
+    })
+    await setUpWhitelist({ email: "@expired.sg", expiry: oneYearAgo })
+    await setUpWhitelist({
+      email: "[email protected]",
+      expiry: oneYearAgo,
+    })
+  })
+
+  it("should show email as whitelisted if the exact email address is whitelisted and expiry is NULL", async () => {
+    // Arrange
+    const email = "[email protected]"
+
+    // Act
+    const result = await isEmailWhitelisted(email)
+
+    // Assert
+    expect(result).toBe(true)
+  })
+
+  it("should show email as whitelisted if the exact email address is whitelisted and expiry is in the future", async () => {
+    // Arrange
+    const email = "[email protected]"
+
+    // Act
+    const result = await isEmailWhitelisted(email)
+
+    // Assert
+    expect(result).toBe(true)
+  })
+
+  it("should show email as not whitelisted if the exact email address is whitelisted and expiry is in the past", async () => {
+    // Arrange
+    const email = "[email protected]"
+
+    // Act
+    const result = await isEmailWhitelisted(email)
+
+    // Assert
+    expect(result).toBe(false)
+  })
+
+  it("should show email as whitelisted if the exact email domain is whitelisted and expiry is NULL", async () => {
+    // Arrange
+    const email = "[email protected]"
+
+    // Act
+    const result = await isEmailWhitelisted(email)
+
+    // Assert
+    expect(result).toBe(true)
+  })
+
+  it("should show email as whitelisted if the exact email domain is whitelisted and expiry is in the future", async () => {
+    // Arrange
+    const email = "[email protected]"
+
+    // Act
+    const result = await isEmailWhitelisted(email)
+
+    // Assert
+    expect(result).toBe(true)
+  })
+
+  it("should show email as not whitelisted if the exact email domain is whitelisted and expiry is in the past", async () => {
+    // Arrange
+    const email = "[email protected]"
+
+    // Act
+    const result = await isEmailWhitelisted(email)
+
+    // Assert
+    expect(result).toBe(false)
+  })
+
+  it("should show email as whitelisted if the suffix of the email domain is whitelisted and expiry is NULL", async () => {
+    // Arrange
+    const email = "[email protected]"
+
+    // Act
+    const result = await isEmailWhitelisted(email)
+
+    // Assert
+    expect(result).toBe(true)
+  })
+
+  it("should show email as whitelisted if the exact email address is expired, but the domain's expiry is in the future", async () => {
+    // Arrange
+    const email = "[email protected]"
+
+    // Act
+    const result = await isEmailWhitelisted(email)
+
+    // Assert
+    expect(result).toBe(true)
+  })
+})

apps/studio/src/server/modules/whitelist/whitelist.service.ts

@@ -0,0 +1,75 @@
+import { TRPCError } from "@trpc/server"
+
+import { isValidEmail } from "~/utils/email"
+import { db } from "../database"
+
+export const isEmailWhitelisted = async (email: string) => {
+  const lowercaseEmail = email.toLowerCase()
+
+  // Extra guard even if Zod validation has already checked
+  if (!isValidEmail(lowercaseEmail)) {
+    throw new TRPCError({
+      code: "BAD_REQUEST",
+      message: "Please sign in with a valid email address.",
+    })
+  }
+
+  // Step 1: Check if the exact email address is whitelisted
+  const exactMatch = await db
+    .selectFrom("Whitelist")
+    .where("email", "=", lowercaseEmail)
+    .where(({ eb }) =>
+      eb.or([eb("expiry", "is", null), eb("expiry", ">", new Date())]),
+    )
+    .select(["id"])
+    .executeTakeFirst()
+
+  if (exactMatch) {
+    return true
+  }
+
+  // Step 2: Check if the exact email domain is whitelisted
+  const emailParts = lowercaseEmail.split("@")
+  if (emailParts.length !== 2) {
+    throw new TRPCError({
+      code: "BAD_REQUEST",
+      message: "Please sign in with a valid email address.",
+    })
+  }
+
+  const emailDomain = `@${emailParts.pop()}`
+  const domainMatch = await db
+    .selectFrom("Whitelist")
+    .where("email", "=", emailDomain)
+    .where(({ eb }) =>
+      eb.or([eb("expiry", "is", null), eb("expiry", ">", new Date())]),
+    )
+    .select(["id"])
+    .executeTakeFirst()
+
+  if (domainMatch) {
+    return true
+  }
+
+  // Step 3: Check if the suffix of the email domain is whitelisted
+  const domainParts = emailDomain.split(".")
+  for (let i = 1; i < domainParts.length; i++) {
+    // Suffices should start with a dot (e.g. ".gov.sg")
+    const suffix = `.${domainParts.slice(i).join(".")}`
+
+    const suffixMatch = await db
+      .selectFrom("Whitelist")
+      .where("email", "=", suffix)
+      .where(({ eb }) =>
+        eb.or([eb("expiry", "is", null), eb("expiry", ">", new Date())]),
+      )
+      .select(["id"])
+      .executeTakeFirst()
+
+    if (suffixMatch) {
+      return true
+    }
+  }
+
+  return false
+}

apps/studio/src/server/trpc.ts

@@ -20,6 +20,7 @@ import getIP from "~/utils/getClientIp"
 import { type Context } from "./context"
 import { defaultMeSelect } from "./modules/me/me.select"
 import { checkRateLimit } from "./modules/rate-limit/rate-limit.service"
+import { isEmailWhitelisted } from "./modules/whitelist/whitelist.service"
 import { prisma } from "./prisma"
 
 interface Meta {
@@ -124,13 +125,6 @@ const baseMiddleware = t.middleware(async ({ ctx, next }) => {
 })
 
 const authMiddleware = t.middleware(async ({ next, ctx }) => {
-  const defaultWhitelist: string[] = []
-  const whitelistedUsers = ctx.gb
-    .getFeatureValue("whitelisted_users", {
-      whitelist: defaultWhitelist,
-    })
-    .whitelist.map((email) => email.toLowerCase())
-
   if (!ctx.session?.userId) {
     throw new TRPCError({ code: "UNAUTHORIZED" })
   }
@@ -145,11 +139,10 @@ const authMiddleware = t.middleware(async ({ next, ctx }) => {
     throw new TRPCError({ code: "UNAUTHORIZED" })
   }
 
-  // check against Growthbook if user is whitelisted for prod/stg
-  if (env.NODE_ENV === "production") {
-    if (!whitelistedUsers.includes(user.email.toLowerCase())) {
-      throw new TRPCError({ code: "UNAUTHORIZED" })
-    }
+  // Ensure that the user is whitelisted to use the app
+  const isWhitelisted = await isEmailWhitelisted(user.email)
+  if (!isWhitelisted) {
+    throw new TRPCError({ code: "UNAUTHORIZED" })
   }
 
   return next({

apps/studio/tests/integration/helpers/auth.ts

@@ -3,8 +3,11 @@ import cuid2 from "@paralleldrive/cuid2"
 import { db } from "~server/db"
 
 import type { User } from "~server/db"
+import { setUpWhitelist } from "./seed"
 
 export const auth = async ({ id, ...user }: SetOptional<User, "id">) => {
+  await setUpWhitelist({ email: user.email })
+
   if (id !== undefined) {
     return db
       .updateTable("User")