refactor checkToken

6 files changed, 80 insertions, 130 deletions

diff --git a/src/api/middlewares/Authentication.ts b/src/api/middlewares/Authentication.ts
index d0e4d8a0..812888a3 100644
--- a/src/api/middlewares/Authentication.ts
+++ b/src/api/middlewares/Authentication.ts
@@ -92,12 +92,7 @@ export async function Authentication(
 	Sentry.setUser({ id: req.user_id });
 
 	try {
-		const { jwtSecret } = Config.get().security;
-
-		const { decoded, user } = await checkToken(
-			req.headers.authorization,
-			jwtSecret,
-		);
+		const { decoded, user } = await checkToken(req.headers.authorization);
 
 		req.token = decoded;
 		req.user_id = decoded.id;
diff --git a/src/api/routes/auth/reset.ts b/src/api/routes/auth/reset.ts
index f97045a6..cb4f8180 100644
--- a/src/api/routes/auth/reset.ts
+++ b/src/api/routes/auth/reset.ts
@@ -48,11 +48,9 @@ router.post(
 	async (req: Request, res: Response) => {
 		const { password, token } = req.body as PasswordResetSchema;
 
-		const { jwtSecret } = Config.get().security;
-
 		let user;
 		try {
-			const userTokenData = await checkToken(token, jwtSecret, true);
+			const userTokenData = await checkToken(token);
 			user = userTokenData.user;
 		} catch {
 			throw FieldErrors({
diff --git a/src/api/routes/auth/verify/index.ts b/src/api/routes/auth/verify/index.ts
index a98c17fa..49f74277 100644
--- a/src/api/routes/auth/verify/index.ts
+++ b/src/api/routes/auth/verify/index.ts
@@ -78,11 +78,10 @@ router.post(
 			}
 		}
 
-		const { jwtSecret } = Config.get().security;
 		let user;
 
 		try {
-			const userTokenData = await checkToken(token, jwtSecret, true);
+			const userTokenData = await checkToken(token);
 			user = userTokenData.user;
 		} catch {
 			throw FieldErrors({
diff --git a/src/gateway/opcodes/Identify.ts b/src/gateway/opcodes/Identify.ts
index 5816c308..837ae351 100644
--- a/src/gateway/opcodes/Identify.ts
+++ b/src/gateway/opcodes/Identify.ts
@@ -30,7 +30,6 @@ import {
 	Intents,
 	Member,
 	ReadyEventData,
-	User,
 	Session,
 	EVENTEnum,
 	Config,
@@ -60,17 +59,6 @@ import { check } from "./instanceOf";
 // TODO: user sharding
 // TODO: check privileged intents, if defined in the config
 
-const getUserFromToken = async (token: string): Promise<string | null> => {
-	try {
-		const { jwtSecret } = Config.get().security;
-		const { decoded } = await checkToken(token, jwtSecret);
-		return decoded.id;
-	} catch (e) {
-		console.error(`[Gateway] Invalid token`, e);
-		return null;
-	}
-};
-
 export async function onIdentify(this: WebSocket, data: Payload) {
 	if (this.user_id) {
 		// we've already identified
@@ -85,12 +73,12 @@ export async function onIdentify(this: WebSocket, data: Payload) {
 
 	this.capabilities = new Capabilities(identify.capabilities || 0);
 
-	// Check auth
-	// TODO: the checkToken call will fetch user, and then we have to refetch with different select
-	// checkToken should be able to select what we want
-	const user_id = await getUserFromToken(identify.token);
-	if (!user_id) return this.close(CLOSECODES.Authentication_failed);
-	this.user_id = user_id;
+	const { user } = await checkToken(identify.token, {
+		relations: ["relationships", "relationships.to", "settings"],
+		select: [...PrivateUserProjection, "relationships"],
+	});
+	if (!user) return this.close(CLOSECODES.Authentication_failed);
+	this.user_id = user.id;
 
 	// Check intents
 	if (!identify.intents) identify.intents = 30064771071n; // TODO: what is this number?
@@ -112,7 +100,7 @@ export async function onIdentify(this: WebSocket, data: Payload) {
 		) {
 			// TODO: why do we even care about this right now?
 			console.log(
-				`[Gateway] Invalid sharding from ${user_id}: ${identify.shard}`,
+				`[Gateway] Invalid sharding from ${user.id}: ${identify.shard}`,
 			);
 			return this.close(CLOSECODES.Invalid_shard);
 		}
@@ -132,22 +120,14 @@ export async function onIdentify(this: WebSocket, data: Payload) {
 	});
 
 	// Get from database:
-	// * the current user,
 	// * the users read states
 	// * guild members for this user
 	// * recipients ( dm channels )
 	// * the bot application, if it exists
-	const [, user, application, read_states, members, recipients] =
-		await Promise.all([
+	const [, application, read_states, members, recipients] = await Promise.all(
+		[
 			session.save(),
 
-			// TODO: Refactor checkToken to allow us to skip this additional query
-			User.findOneOrFail({
-				where: { id: this.user_id },
-				relations: ["relationships", "relationships.to", "settings"],
-				select: [...PrivateUserProjection, "relationships"],
-			}),
-
 			Application.findOne({
 				where: { id: this.user_id },
 				select: ["id", "flags"],
@@ -224,7 +204,8 @@ export async function onIdentify(this: WebSocket, data: Payload) {
 					},
 				},
 			}),
-		]);
+		],
+	);
 
 	// We forgot to migrate user settings from the JSON column of `users`
 	// to the `user_settings` table theyre in now,
@@ -407,6 +388,17 @@ export async function onIdentify(this: WebSocket, data: Payload) {
 		merged_members: merged_members,
 		sessions: allSessions,
 
+		resume_gateway_url:
+			Config.get().gateway.endpointClient ||
+			Config.get().gateway.endpointPublic ||
+			"ws://127.0.0.1:3001",
+
+		// lol hack whatever
+		required_action:
+			Config.get().login.requireVerification && !user.verified
+				? "REQUIRE_VERIFIED_EMAIL"
+				: undefined,
+
 		consents: {
 			personalization: {
 				consented: false, // TODO
@@ -421,18 +413,8 @@ export async function onIdentify(this: WebSocket, data: Payload) {
 		friend_suggestion_count: 0,
 		analytics_token: "",
 		tutorial: null,
-		resume_gateway_url:
-			Config.get().gateway.endpointClient ||
-			Config.get().gateway.endpointPublic ||
-			"ws://127.0.0.1:3001",
 		session_type: "normal", // TODO
 		auth_session_id_hash: "", // TODO
-
-		// lol hack whatever
-		required_action:
-			Config.get().login.requireVerification && !user.verified
-				? "REQUIRE_VERIFIED_EMAIL"
-				: undefined,
 	};
 
 	// Send READY
diff --git a/src/util/schemas/RegisterSchema.ts b/src/util/schemas/RegisterSchema.ts
index f6c99b18..7b7de9c7 100644
--- a/src/util/schemas/RegisterSchema.ts
+++ b/src/util/schemas/RegisterSchema.ts
@@ -42,4 +42,8 @@ export interface RegisterSchema {
 	captcha_key?: string;
 
 	promotional_email_opt_in?: boolean;
+
+	// part of pomelo
+	unique_username_registration?: boolean;
+	global_name?: string;
 }
diff --git a/src/util/util/Token.ts b/src/util/util/Token.ts
index 90310176..eec72522 100644
--- a/src/util/util/Token.ts
+++ b/src/util/util/Token.ts
@@ -19,94 +19,66 @@
 import jwt, { VerifyOptions } from "jsonwebtoken";
 import { Config } from "./Config";
 import { User } from "../entities";
+// TODO: dont use deprecated APIs lol
+import {
+	FindOptionsRelationByString,
+	FindOptionsSelectByString,
+} from "typeorm";
 
 export const JWTOptions: VerifyOptions = { algorithms: ["HS256"] };
 
 export type UserTokenData = {
 	user: User;
-	decoded: { id: string; iat: number };
+	decoded: { id: string; iat: number; email?: string };
 };
 
-async function checkEmailToken(
-	decoded: jwt.JwtPayload,
-): Promise<UserTokenData> {
-	// eslint-disable-next-line no-async-promise-executor
-	return new Promise(async (res, rej) => {
-		if (!decoded.iat) return rej("Invalid Token"); // will never happen, just for typings.
-
-		const user = await User.findOne({
-			where: {
-				email: decoded.email,
-			},
-			select: [
-				"email",
-				"id",
-				"verified",
-				"deleted",
-				"disabled",
-				"username",
-				"data",
-			],
-		});
-
-		if (!user) return rej("Invalid Token");
-
-		if (new Date().getTime() > decoded.iat * 1000 + 86400 * 1000)
-			return rej("Invalid Token");
-
-		// Using as here because we assert `id` and `iat` are in decoded.
-		// TS just doesn't want to assume its there, though.
-		return res({ decoded, user } as UserTokenData);
-	});
-}
-
-export function checkToken(
+export const checkToken = (
 	token: string,
-	jwtSecret: string,
-	isEmailVerification = false,
-): Promise<UserTokenData> {
-	return new Promise((res, rej) => {
-		token = token.replace("Bot ", "");
-		token = token.replace("Bearer ", "");
-		/**
-		in spacebar, even with instances that have bot distinction; we won't enforce "Bot" prefix,
-		as we don't really have separate pathways for bots 
-		**/
-
-		jwt.verify(token, jwtSecret, JWTOptions, async (err, decoded) => {
-			if (err || !decoded) return rej("Invalid Token");
-			if (
-				typeof decoded == "string" ||
-				!("id" in decoded) ||
-				!decoded.iat
-			)
-				return rej("Invalid Token"); // will never happen, just for typings.
-
-			if (isEmailVerification) return res(checkEmailToken(decoded));
-
-			const user = await User.findOne({
-				where: { id: decoded.id },
-				select: ["data", "bot", "disabled", "deleted", "rights"],
-			});
-
-			if (!user) return rej("Invalid Token");
-
-			// we need to round it to seconds as it saved as seconds in jwt iat and valid_tokens_since is stored in milliseconds
-			if (
-				decoded.iat * 1000 <
-				new Date(user.data.valid_tokens_since).setSeconds(0, 0)
-			)
-				return rej("Invalid Token");
-
-			if (user.disabled) return rej("User disabled");
-			if (user.deleted) return rej("User not found");
-
-			// Using as here because we assert `id` and `iat` are in decoded.
-			// TS just doesn't want to assume its there, though.
-			return res({ decoded, user } as UserTokenData);
-		});
+	opts?: {
+		select?: FindOptionsSelectByString<User>;
+		relations?: FindOptionsRelationByString;
+	},
+): Promise<UserTokenData> =>
+	new Promise((resolve, reject) => {
+		jwt.verify(
+			token,
+			Config.get().security.jwtSecret,
+			JWTOptions,
+			async (err, out) => {
+				const decoded = out as UserTokenData["decoded"];
+				if (err || !decoded) return reject("Invalid Token");
+
+				const user = await User.findOne({
+					where: decoded.email
+						? { email: decoded.email }
+						: { id: decoded.id },
+					select: [
+						...(opts?.select || []),
+						"bot",
+						"disabled",
+						"deleted",
+						"rights",
+						"data",
+					],
+					relations: opts?.relations,
+				});
+
+				if (!user) return reject("User not found");
+
+				// we need to round it to seconds as it saved as seconds in jwt iat and valid_tokens_since is stored in milliseconds
+				if (
+					decoded.iat * 1000 <
+					new Date(user.data.valid_tokens_since).setSeconds(0, 0)
+				)
+					return reject("Invalid Token");
+
+				if (user.disabled) return reject("User disabled");
+				if (user.deleted) return reject("User not found");
+
+				return resolve({ decoded, user });
+			},
+		);
 	});
-}
 
 export async function generateToken(id: string, email?: string) {
 	const iat = Math.floor(Date.now() / 1000);