v4.5.1 (#417)

packages/service/support/user/auth.ts

@@ -0,0 +1,206 @@
+import type { NextApiResponse, NextApiRequest } from 'next';
+import Cookie from 'cookie';
+import jwt from 'jsonwebtoken';
+import { authOpenApiKey } from '../openapi/auth';
+import { authOutLinkId } from '../outLink/auth';
+import { MongoUser } from './schema';
+import type { UserModelSchema } from '@fastgpt/global/support/user/type';
+import { ERROR_ENUM } from '@fastgpt/global/common/error/errorCode';
+
+export enum AuthUserTypeEnum {
+  token = 'token',
+  root = 'root',
+  apikey = 'apikey',
+  outLink = 'outLink'
+}
+
+/* auth balance */
+export const authBalanceByUid = async (uid: string) => {
+  const user = await MongoUser.findById<UserModelSchema>(
+    uid,
+    '_id username balance openaiAccount timezone'
+  );
+  if (!user) {
+    return Promise.reject(ERROR_ENUM.unAuthorization);
+  }
+
+  if (user.balance <= 0) {
+    return Promise.reject(ERROR_ENUM.insufficientQuota);
+  }
+  return user;
+};
+
+/* uniform auth user */
+export const authUser = async ({
+  req,
+  authToken = false,
+  authRoot = false,
+  authApiKey = false,
+  authBalance = false,
+  authOutLink
+}: {
+  req: NextApiRequest;
+  authToken?: boolean;
+  authRoot?: boolean;
+  authApiKey?: boolean;
+  authBalance?: boolean;
+  authOutLink?: boolean;
+}) => {
+  const authCookieToken = async (cookie?: string, token?: string): Promise<string> => {
+    // 获取 cookie
+    const cookies = Cookie.parse(cookie || '');
+    const cookieToken = cookies.token || token;
+
+    if (!cookieToken) {
+      return Promise.reject(ERROR_ENUM.unAuthorization);
+    }
+
+    return await authJWT(cookieToken);
+  };
+  // from authorization get apikey
+  const parseAuthorization = async (authorization?: string) => {
+    if (!authorization) {
+      return Promise.reject(ERROR_ENUM.unAuthorization);
+    }
+
+    // Bearer fastgpt-xxxx-appId
+    const auth = authorization.split(' ')[1];
+    if (!auth) {
+      return Promise.reject(ERROR_ENUM.unAuthorization);
+    }
+
+    const { apikey, appId: authorizationAppid = '' } = await (async () => {
+      const arr = auth.split('-');
+      // abandon
+      if (arr.length === 3) {
+        return {
+          apikey: `${arr[0]}-${arr[1]}`,
+          appId: arr[2]
+        };
+      }
+      if (arr.length === 2) {
+        return {
+          apikey: auth
+        };
+      }
+      return Promise.reject(ERROR_ENUM.unAuthorization);
+    })();
+
+    // auth apikey
+    const { userId, appId: apiKeyAppId = '' } = await authOpenApiKey({ apikey });
+
+    return {
+      uid: userId,
+      apikey,
+      appId: apiKeyAppId || authorizationAppid
+    };
+  };
+  // root user
+  const parseRootKey = async (rootKey?: string, userId = '') => {
+    if (!rootKey || !process.env.ROOT_KEY || rootKey !== process.env.ROOT_KEY) {
+      return Promise.reject(ERROR_ENUM.unAuthorization);
+    }
+    return userId;
+  };
+
+  const { cookie, token, apikey, rootkey, userid, authorization } = (req.headers || {}) as {
+    cookie?: string;
+    token?: string;
+    apikey?: string;
+    rootkey?: string; // abandon
+    userid?: string;
+    authorization?: string;
+  };
+  const { shareId } = (req?.body || {}) as { shareId?: string };
+
+  let uid = '';
+  let appId = '';
+  let openApiKey = apikey;
+  let authType: `${AuthUserTypeEnum}` = AuthUserTypeEnum.token;
+
+  if (authOutLink && shareId) {
+    const res = await authOutLinkId({ id: shareId });
+    uid = res.userId;
+    authType = AuthUserTypeEnum.outLink;
+  } else if (authToken && (cookie || token)) {
+    // user token(from fastgpt web)
+    uid = await authCookieToken(cookie, token);
+    authType = AuthUserTypeEnum.token;
+  } else if (authRoot && rootkey) {
+    // root user
+    uid = await parseRootKey(rootkey, userid);
+    authType = AuthUserTypeEnum.root;
+  } else if (authApiKey && apikey) {
+    // apikey
+    const parseResult = await authOpenApiKey({ apikey });
+    uid = parseResult.userId;
+    authType = AuthUserTypeEnum.apikey;
+    openApiKey = parseResult.apikey;
+  } else if (authApiKey && authorization) {
+    // apikey from authorization
+    const authResponse = await parseAuthorization(authorization);
+    uid = authResponse.uid;
+    appId = authResponse.appId;
+    openApiKey = authResponse.apikey;
+    authType = AuthUserTypeEnum.apikey;
+  }
+
+  // not rootUser and no uid, reject request
+  if (!rootkey && !uid) {
+    return Promise.reject(ERROR_ENUM.unAuthorization);
+  }
+
+  // balance check
+  const user = await (() => {
+    if (authBalance) {
+      return authBalanceByUid(uid);
+    }
+  })();
+
+  return {
+    userId: String(uid),
+    appId,
+    authType,
+    user,
+    apikey: openApiKey
+  };
+};
+
+/* 生成 token */
+export function generateToken(userId: string) {
+  const key = process.env.TOKEN_KEY as string;
+  const token = jwt.sign(
+    {
+      userId,
+      exp: Math.floor(Date.now() / 1000) + 60 * 60 * 24 * 7
+    },
+    key
+  );
+  return token;
+}
+// auth token
+export function authJWT(token: string) {
+  return new Promise<string>((resolve, reject) => {
+    const key = process.env.TOKEN_KEY as string;
+
+    jwt.verify(token, key, function (err, decoded: any) {
+      if (err || !decoded?.userId) {
+        reject(ERROR_ENUM.unAuthorization);
+        return;
+      }
+      resolve(decoded.userId);
+    });
+  });
+}
+
+/* set cookie */
+export const setCookie = (res: NextApiResponse, token: string) => {
+  res.setHeader(
+    'Set-Cookie',
+    `token=${token}; Path=/; HttpOnly; Max-Age=604800; Samesite=None; Secure;`
+  );
+};
+/* clear cookie */
+export const clearCookie = (res: NextApiResponse) => {
+  res.setHeader('Set-Cookie', 'token=; Path=/; Max-Age=0');
+};

packages/service/support/user/schema.ts

@@ -0,0 +1,63 @@
+import { connectionMongo, type Model } from '../../common/mongo';
+const { Schema, model, models } = connectionMongo;
+import { hashStr } from '@fastgpt/global/common/string/tools';
+import { PRICE_SCALE } from '@fastgpt/global/common/bill/constants';
+import type { UserModelSchema } from '@fastgpt/global/support/user/type';
+
+const UserSchema = new Schema({
+  username: {
+    // 可以是手机/邮箱，新的验证都只用手机
+    type: String,
+    required: true,
+    unique: true // 唯一
+  },
+  password: {
+    type: String,
+    required: true,
+    set: (val: string) => hashStr(val),
+    get: (val: string) => hashStr(val),
+    select: false
+  },
+  createTime: {
+    type: Date,
+    default: () => new Date()
+  },
+  avatar: {
+    type: String,
+    default: '/icon/human.svg'
+  },
+  balance: {
+    type: Number,
+    default: 2 * PRICE_SCALE
+  },
+  inviterId: {
+    // 谁邀请注册的
+    type: Schema.Types.ObjectId,
+    ref: 'user'
+  },
+  promotionRate: {
+    type: Number,
+    default: 15
+  },
+  limit: {
+    exportKbTime: {
+      // Every half hour
+      type: Date
+    },
+    datasetMaxCount: {
+      type: Number
+    }
+  },
+  openaiAccount: {
+    type: {
+      key: String,
+      baseUrl: String
+    }
+  },
+  timezone: {
+    type: String,
+    default: 'Asia/Shanghai'
+  }
+});
+
+export const MongoUser: Model<UserModelSchema> = models['user'] || model('user', UserSchema);