From f86d944c25735b598b7d974c754619141e1e90c7 Mon Sep 17 00:00:00 2001
From: beilunyang <786220806@qq.com>
Date: Thu, 27 Feb 2025 23:59:36 +0800
Subject: [PATCH] feat: Add role-based email limit exemption for emperors

---
 app/api/emails/generate/route.ts     | 34 ++++++++++++++++------------
 app/components/emails/email-list.tsx |  9 +++++++-
 app/hooks/use-user-role.ts           | 21 +++++++++++++++++
 app/lib/auth.ts                      |  9 ++++++++
 4 files changed, 58 insertions(+), 15 deletions(-)
 create mode 100644 app/hooks/use-user-role.ts

diff --git a/app/api/emails/generate/route.ts b/app/api/emails/generate/route.ts
index 558553e..2cea299 100644
--- a/app/api/emails/generate/route.ts
+++ b/app/api/emails/generate/route.ts
@@ -7,29 +7,35 @@ import { EXPIRY_OPTIONS } from "@/types/email"
 import { EMAIL_CONFIG } from "@/config"
 import { getRequestContext } from "@cloudflare/next-on-pages"
 import { getUserId } from "@/lib/apiKey"
+import { getUserRole } from "@/lib/auth"
+import { ROLES } from "@/lib/permissions"
+
 export const runtime = "edge"
 
 export async function POST(request: Request) {
   const db = createDb()
 
   const userId = await getUserId()
+  const userRole = await getUserRole(userId!)
 
   try {
-    const activeEmailsCount = await db
-      .select({ count: sql<number>`count(*)` })
-      .from(emails)
-      .where(
-        and(
-          eq(emails.userId, userId!),
-          gt(emails.expiresAt, new Date())
+    if (userRole !== ROLES.EMPEROR) {
+      const activeEmailsCount = await db
+        .select({ count: sql<number>`count(*)` })
+        .from(emails)
+        .where(
+          and(
+            eq(emails.userId, userId!),
+            gt(emails.expiresAt, new Date())
+          )
         )
-      )
-    
-    if (Number(activeEmailsCount[0].count) >= EMAIL_CONFIG.MAX_ACTIVE_EMAILS) {
-      return NextResponse.json(
-        { error: `已达到最大邮箱数量限制 (${EMAIL_CONFIG.MAX_ACTIVE_EMAILS})` },
-        { status: 403 }
-      )
+      
+      if (Number(activeEmailsCount[0].count) >= EMAIL_CONFIG.MAX_ACTIVE_EMAILS) {
+        return NextResponse.json(
+          { error: `已达到最大邮箱数量限制 (${EMAIL_CONFIG.MAX_ACTIVE_EMAILS})` },
+          { status: 403 }
+        )
+      }
     }
 
     const { name, expiryTime, domain } = await request.json<{ 
diff --git a/app/components/emails/email-list.tsx b/app/components/emails/email-list.tsx
index 71ee3e6..b304538 100644
--- a/app/components/emails/email-list.tsx
+++ b/app/components/emails/email-list.tsx
@@ -19,6 +19,8 @@ import {
   AlertDialogHeader,
   AlertDialogTitle,
 } from "@/components/ui/alert-dialog"
+import { ROLES } from "@/lib/permissions"
+import { useUserRole } from "@/hooks/use-user-role"
 
 interface Email {
   id: string
@@ -40,6 +42,7 @@ interface EmailResponse {
 
 export function EmailList({ onEmailSelect, selectedEmailId }: EmailListProps) {
   const { data: session } = useSession()
+  const { role } = useUserRole()
   const [emails, setEmails] = useState<Email[]>([])
   const [loading, setLoading] = useState(true)
   const [refreshing, setRefreshing] = useState(false)
@@ -167,7 +170,11 @@ export function EmailList({ onEmailSelect, selectedEmailId }: EmailListProps) {
               <RefreshCw className="h-4 w-4" />
             </Button>
             <span className="text-xs text-gray-500">
-              {total}/{EMAIL_CONFIG.MAX_ACTIVE_EMAILS} 个邮箱
+              {role === ROLES.EMPEROR ? (
+                `${total}/∞ 个邮箱`
+              ) : (
+                `${total}/${EMAIL_CONFIG.MAX_ACTIVE_EMAILS} 个邮箱`
+              )}
             </span>
           </div>
           <CreateDialog onEmailCreated={handleRefresh} />
diff --git a/app/hooks/use-user-role.ts b/app/hooks/use-user-role.ts
new file mode 100644
index 0000000..6bd7f0e
--- /dev/null
+++ b/app/hooks/use-user-role.ts
@@ -0,0 +1,21 @@
+"use client"
+
+import { useSession } from "next-auth/react"
+import { Role } from "@/lib/permissions"
+import { useEffect, useState } from "react"
+
+export function useUserRole() {
+  const { data: session } = useSession()
+  const [role, setRole] = useState<Role | null>(null)
+
+  useEffect(() => {
+    if (session?.user?.roles?.[0]?.name) {
+      setRole(session.user.roles[0].name as Role)
+    }
+  }, [session])
+
+  return {
+    role,
+    loading: !session
+  }
+} 
\ No newline at end of file
diff --git a/app/lib/auth.ts b/app/lib/auth.ts
index b715fc7..c1f48cc 100644
--- a/app/lib/auth.ts
+++ b/app/lib/auth.ts
@@ -52,6 +52,15 @@ export async function assignRoleToUser(db: Db, userId: string, roleId: string) {
     })
 }
 
+export async function getUserRole(userId: string) {
+  const db = createDb()
+  const userRoleRecords = await db.query.userRoles.findMany({
+    where: eq(userRoles.userId, userId),
+    with: { role: true },
+  })
+  return userRoleRecords[0].role.name
+}
+
 export async function checkPermission(permission: Permission) {
   const session = await auth()
   if (!session?.user?.id) return false