moemail/app/lib/auth.ts

import NextAuth from "next-auth"
import GitHub from "next-auth/providers/github"
import { DrizzleAdapter } from "@auth/drizzle-adapter"
import { createDb, Db } from "./db"
import { accounts, sessions, users, roles, userRoles } from "./schema"
import { eq } from "drizzle-orm"
import { Permission, hasPermission, ROLES, Role } from "./permissions"

const ROLE_DESCRIPTIONS: Record<Role, string> = {
  [ROLES.EMPEROR]: "皇帝（网站所有者）",
  [ROLES.KNIGHT]: "骑士（高级用户）",
  [ROLES.CIVILIAN]: "平民（普通用户）",
}

const getDefaultRole = (): Role =>
  process.env.OPEN_REGISTRATION === 'true' ? ROLES.KNIGHT : ROLES.CIVILIAN

async function findOrCreateRole(db: Db, roleName: Role) {
  let role = await db.query.roles.findFirst({
    where: eq(roles.name, roleName),
  })

  if (!role) {
    const [newRole] = await db.insert(roles)
      .values({
        name: roleName,
        description: ROLE_DESCRIPTIONS[roleName],
      })
      .returning()
    role = newRole
  }

  return role
}

async function assignRoleToUser(db: Db, userId: string, roleId: string) {
  await db.insert(userRoles)
    .values({
      userId,
      roleId,
    })
}

export async function checkPermission(permission: Permission) {
  const session = await auth()
  if (!session?.user?.id) return false

  const db = createDb()
  const userRoleRecords = await db.query.userRoles.findMany({
    where: eq(userRoles.userId, session.user.id),
    with: { role: true },
  })

  const userRoleNames = userRoleRecords.map(ur => ur.role.name)
  return hasPermission(userRoleNames as Role[], permission)
}

export const {
  handlers: { GET, POST },
  auth,
  signIn,
  signOut
} = NextAuth(() => ({
  secret: process.env.AUTH_SECRET,
  adapter: DrizzleAdapter(createDb(), {
    usersTable: users,
    accountsTable: accounts,
    sessionsTable: sessions,
  }),
  providers: [
    GitHub({
      clientId: process.env.AUTH_GITHUB_ID,
      clientSecret: process.env.AUTH_GITHUB_SECRET,
    })
  ],
  events: {
    async signIn({ user }) {
      if (!user.id) return

      try {
        const db = createDb()
        const existingRole = await db.query.userRoles.findFirst({
          where: eq(userRoles.userId, user.id),
        })

        if (existingRole) return

        const defaultRole = await findOrCreateRole(db, getDefaultRole())
        await assignRoleToUser(db, user.id, defaultRole.id)
      } catch (error) {
        console.error('Error assigning role:', error)
      }
    },
  },
  pages: {
    signIn: "/",
    error: "/",
  },
  callbacks: {
    async session({ session, user }) {
      if (!session?.user) return session

      const db = createDb()
      let userRoleRecords = await db.query.userRoles.findMany({
        where: eq(userRoles.userId, user.id),
        with: { role: true },
      })

      if (!userRoleRecords.length) {
        const defaultRole = await findOrCreateRole(db, getDefaultRole())
        await assignRoleToUser(db, user.id, defaultRole.id)
        userRoleRecords = [{
          userId: user.id,
          roleId: defaultRole.id,
          createdAt: new Date(),
          role: defaultRole
        }]
      }

      session.user.roles = userRoleRecords.map(ur => ({
        name: ur.role.name,
      }))

      return session
    },
  }
}))