chore: set user role

fix: register failed (Close #160)

Author: Kerwin

README.md

@@ -417,11 +417,11 @@ A: 一种可能原因是经过 Nginx 反向代理，开启了 buffer，则 Nginx
 
 <div style="display: flex; gap: 20px;">
 	<div style="text-align: center">
-		<img style="max-width: 100%" src="./docs/wechat.png" alt="微信" />
+		<img style="max-width: 200px" src="./docs/wechat.png" alt="微信" />
 		<p>WeChat Pay</p>
 	</div>
 	<div style="text-align: center">
-		<img style="max-width: 100%" src="./docs/alipay.png" alt="支付宝" />
+		<img style="max-width: 200px" src="./docs/alipay.png" alt="支付宝" />
 		<p>Alipay</p>
 	</div>
 </div>

service/src/index.ts

@@ -5,10 +5,10 @@ import { ObjectId } from 'mongodb'
 import type { RequestProps } from './types'
 import type { ChatContext, ChatMessage } from './chatgpt'
 import { chatConfig, chatReplyProcess, containsSensitiveWords, getRandomApiKey, initAuditService } from './chatgpt'
-import { auth } from './middleware/auth'
-import { clearConfigCache, getApiKeys, getCacheConfig, getOriginConfig } from './storage/config'
+import { auth, getUserId } from './middleware/auth'
+import { clearApiKeyCache, clearConfigCache, getApiKeys, getCacheApiKeys, getCacheConfig, getOriginConfig } from './storage/config'
 import type { AuditConfig, CHATMODEL, ChatInfo, ChatOptions, Config, KeyConfig, MailConfig, SiteConfig, UsageResponse, UserInfo } from './storage/model'
-import { Status, UserRole } from './storage/model'
+import { Status, UserRole, chatModelOptions } from './storage/model'
 import {
   clearChat,
   createChatRoom,
@@ -36,12 +36,13 @@ import {
   updateUserChatModel,
   updateUserInfo,
   updateUserPassword,
+  updateUserRole,
   updateUserStatus,
   upsertKey,
   verifyUser,
 } from './storage/mongo'
 import { limiter } from './middleware/limiter'
-import { isEmail, isNotEmptyString } from './utils/is'
+import { hasAnyRole, isEmail, isNotEmptyString } from './utils/is'
 import { sendNoticeMail, sendResetPasswordMail, sendTestMail, sendVerifyMail, sendVerifyMailAdmin } from './utils/mail'
 import { checkUserResetPassword, checkUserVerify, checkUserVerifyAdmin, getUserResetPasswordUrl, getUserVerifyUrl, getUserVerifyUrlAdmin, md5 } from './utils/security'
 import { rootAuth } from './middleware/rootAuth'
@@ -558,8 +559,51 @@ router.post('/session', async (req, res) => {
     const allowRegister = (await getCacheConfig()).siteConfig.registerEnabled
     if (config.apiModel !== 'ChatGPTAPI' && config.apiModel !== 'ChatGPTUnofficialProxyAPI')
       config.apiModel = 'ChatGPTAPI'
+    const userId = await getUserId(req)
+    const chatModels: {
+      label
+      key: string
+      value: string
+    }[] = []
+    if (userId != null) {
+      const user = await getUserById(userId)
+      const keys = (await getCacheApiKeys()).filter(d => hasAnyRole(d.userRoles, user.roles))
+
+      const count: { key: string; count: number }[] = []
+      chatModelOptions.forEach((chatModel) => {
+        keys.forEach((key) => {
+          if (key.chatModels.includes(chatModel.value)) {
+            if (count.filter(d => d.key === chatModel.value).length <= 0) {
+              count.push({ key: chatModel.value, count: 1 })
+            }
+            else {
+              const thisCount = count.filter(d => d.key === chatModel.value)[0]
+              thisCount.count++
+            }
+          }
+        })
+      })
+      count.forEach((c) => {
+        const thisChatModel = chatModelOptions.filter(d => d.value === c.key)[0]
+        chatModels.push({
+          label: `${thisChatModel.label} (${c.count})`,
+          key: c.key,
+          value: c.key,
+        })
+      })
+    }
 
-    res.send({ status: 'Success', message: '', data: { auth: hasAuth, allowRegister, model: config.apiModel, title: config.siteConfig.siteTitle } })
+    res.send({
+      status: 'Success',
+      message: '',
+      data: {
+        auth: hasAuth,
+        allowRegister,
+        model: config.apiModel,
+        title: config.siteConfig.siteTitle,
+        chatModels,
+      },
+    })
   }
   catch (error) {
     res.send({ status: 'Fail', message: error.message, data: null })
@@ -693,6 +737,17 @@ router.post('/user-status', rootAuth, async (req, res) => {
   }
 })
 
+router.post('/user-role', rootAuth, async (req, res) => {
+  try {
+    const { userId, roles } = req.body as { userId: string; roles: UserRole[] }
+    await updateUserRole(userId, roles)
+    res.send({ status: 'Success', message: '更新成功 | Update successfully' })
+  }
+  catch (error) {
+    res.send({ status: 'Fail', message: error.message, data: null })
+  }
+})
+
 router.post('/verify', async (req, res) => {
   try {
     const { token } = req.body as { token: string }
@@ -873,6 +928,7 @@ router.post('/setting-key-upsert', rootAuth, async (req, res) => {
     if (keyConfig._id !== undefined)
       keyConfig._id = new ObjectId(keyConfig._id)
     await upsertKey(keyConfig)
+    clearApiKeyCache()
     res.send({ status: 'Success', message: '成功 | Successfully' })
   }
   catch (error) {

service/src/middleware/auth.ts

@@ -1,4 +1,5 @@
 import jwt from 'jsonwebtoken'
+import type { Request } from 'express'
 import { getCacheConfig } from '../storage/config'
 import { getUserById } from '../storage/mongo'
 import { Status } from '../storage/model'
@@ -27,4 +28,17 @@ const auth = async (req, res, next) => {
   }
 }
 
-export { auth }
+async function getUserId(req: Request): Promise<string | undefined> {
+  try {
+    const token = req.header('Authorization').replace('Bearer ', '')
+    const config = await getCacheConfig()
+    const info = jwt.verify(token, config.siteConfig.loginSalt.trim())
+    return info.userId
+  }
+  catch (error) {
+
+  }
+  return null
+}
+
+export { auth, getUserId }

service/src/storage/config.ts

@@ -133,6 +133,11 @@ export async function getCacheApiKeys(): Promise<KeyConfig[]> {
   return Promise.resolve(apiKeysCachedConfig)
 }
 
+export function clearApiKeyCache() {
+  apiKeysCacheExpiration = 0
+  getCacheApiKeys()
+}
+
 export async function getApiKeys() {
   const result = await getKeys()
   if (result.keys.length <= 0) {

service/src/storage/model.ts

@@ -61,6 +61,27 @@ export const CHATMODELS: CHATMODEL[] = [
   'gpt-4-browsing',
 ]
 
+export const chatModelOptions = [
+  'gpt-3.5-turbo',
+  'gpt-3.5-turbo-0301',
+  'gpt-4',
+  'gpt-4-0314',
+  'gpt-4-32k',
+  'gpt-4-32k-0314',
+  'text-davinci-002-render-sha-mobile',
+  'gpt-4-mobile',
+  'gpt-4-browsing',
+].map((model: string) => {
+  let label = model
+  if (model === 'text-davinci-002-render-sha-mobile')
+    label = 'gpt-3.5-mobile'
+  return {
+    label,
+    key: model,
+    value: model,
+  }
+})
+
 export class ChatRoom {
   _id: ObjectId
   roomId: number

service/src/storage/mongo.ts

@@ -232,6 +232,8 @@ export async function getUserById(userId: string): Promise<UserInfo> {
 }
 
 function initUserInfo(userInfo: UserInfo) {
+  if (userInfo == null)
+    return
   if (userInfo.config == null)
     userInfo.config = new UserConfig()
   if (userInfo.config.chatModel == null)
@@ -252,6 +254,10 @@ export async function updateUserStatus(userId: string, status: Status) {
   return await userCol.updateOne({ _id: new ObjectId(userId) }, { $set: { status, verifyTime: new Date().toLocaleString() } })
 }
 
+export async function updateUserRole(userId: string, roles: UserRole[]) {
+  return await userCol.updateOne({ _id: new ObjectId(userId) }, { $set: { roles, verifyTime: new Date().toLocaleString() } })
+}
+
 export async function getConfig(): Promise<Config> {
   return await configCol.findOne() as Config
 }

service/src/utils/is.ts

@@ -38,6 +38,7 @@ export function isTextAudioType(value: any): value is TextAudioType {
     || value === TextAudioType.All
   )
 }
+
 export function hasAnyRole(userRoles: UserRole[] | undefined, roles: UserRole[]): boolean {
   if (!userRoles || userRoles.length === 0 || !roles || roles.length === 0)
     return false

src/api/index.ts

@@ -1,6 +1,6 @@
 import type { AxiosProgressEvent, GenericAbortSignal } from 'axios'
 import { get, post } from '@/utils/request'
-import type { AuditConfig, CHATMODEL, ConfigState, KeyConfig, MailConfig, SiteConfig, Status } from '@/components/common/Setting/model'
+import type { AuditConfig, CHATMODEL, ConfigState, KeyConfig, MailConfig, SiteConfig, Status, UserRole } from '@/components/common/Setting/model'
 import { useAuthStore, useSettingStore } from '@/store'
 
 export function fetchChatAPI<T = any>(
@@ -142,6 +142,13 @@ export function fetchUpdateUserStatus<T = any>(userId: string, status: Status) {
   })
 }
 
+export function fetchUpdateUserRole<T = any>(userId: string, roles: UserRole[]) {
+  return post<T>({
+    url: '/user-role',
+    data: { userId, roles },
+  })
+}
+
 export function fetchGetChatRooms<T = any>() {
   return get<T>({
     url: '/chatrooms',

src/components/common/Setting/Keys.vue

@@ -2,12 +2,14 @@
 import { h, onMounted, reactive, ref } from 'vue'
 import { NButton, NDataTable, NInput, NModal, NSelect, NSpace, NSwitch, NTag, useDialog, useMessage } from 'naive-ui'
 import type { CHATMODEL } from './model'
-import { KeyConfig, Status, UserRole, apiModelOptions, chatModelOptions, userRoleOptions } from './model'
+import { KeyConfig, Status, UserRole, apiModelOptions, userRoleOptions } from './model'
 import { fetchGetKeys, fetchUpdateApiKeyStatus, fetchUpsertApiKey } from '@/api'
 import { t } from '@/locales'
+import { useAuthStore } from '@/store'
 
 const ms = useMessage()
 const dialog = useDialog()
+const authStore = useAuthStore()
 const loading = ref(false)
 const show = ref(false)
 const handleSaving = ref(false)
@@ -254,7 +256,7 @@ onMounted(async () => {
               style="width: 100%"
               multiple
               :value="keyConfig.chatModels"
-              :options="chatModelOptions"
+              :options="authStore.session?.chatModels"
               @update-value="value => keyConfig.chatModels = value"
             />
           </div>