From fa3052d36e5ec9bc5965ca7417e63a3216289d7a Mon Sep 17 00:00:00 2001
From: zkldi <ktchidev@gmail.com>
Date: Mon, 1 Nov 2021 01:22:38 +0000
Subject: [PATCH] Add support for ENABLE_PUBLIC_SIGNUP

---
 server/src/lib/setup/config.ts                |  2 +
 .../src/server/router/api/v1/auth/router.ts   | 53 ++++++++++---------
 2 files changed, 30 insertions(+), 25 deletions(-)

diff --git a/server/src/lib/setup/config.ts b/server/src/lib/setup/config.ts
index 517306d8a..54dc2f45c 100644
--- a/server/src/lib/setup/config.ts
+++ b/server/src/lib/setup/config.ts
@@ -74,6 +74,7 @@ export interface TachiConfig {
 	OUR_URL: string;
 	LOGGER_DISCORD_WEBHOOK?: string;
 	DISCORD_WHO_TO_TAG?: string[];
+	ENABLE_PUBLIC_SIGNUP?: boolean;
 }
 
 const isValidOauth2 = p.optional({
@@ -112,6 +113,7 @@ const err = p(config, {
 	OUR_URL: "string",
 	LOGGER_DISCORD_WEBHOOK: "*string",
 	DISCORD_WHO_TO_TAG: p.optional(["string"]),
+	ENABLE_PUBLIC_SIGNUP: "*boolean",
 });
 
 if (err) {
diff --git a/server/src/server/router/api/v1/auth/router.ts b/server/src/server/router/api/v1/auth/router.ts
index aee7a549b..5269ba247 100644
--- a/server/src/server/router/api/v1/auth/router.ts
+++ b/server/src/server/router/api/v1/auth/router.ts
@@ -161,7 +161,7 @@ router.post(
 			username: Prudence.regex(/^[a-zA-Z_-][a-zA-Z0-9_-]{2,20}$/u),
 			password: ValidatePassword,
 			email: Prudence.regex(LAZY_EMAIL_REGEX),
-			inviteCode: "string",
+			inviteCode: "*string",
 			captcha: "string",
 		},
 		{
@@ -218,30 +218,33 @@ router.post(
 
 		try {
 			const userID = await GetNextCounterValue("users");
-			const inviteCodeDoc = await db.invites.findOneAndUpdate(
-				{
-					code: req.body.inviteCode,
-					consumed: false,
-				},
-				{
-					$set: {
-						consumed: true,
-						consumedAt: Date.now(),
-						consumedBy: userID,
+
+			if (!ServerConfig.ENABLE_PUBLIC_SIGNUP) {
+				const inviteCodeDoc = await db.invites.findOneAndUpdate(
+					{
+						code: req.body.inviteCode,
+						consumed: false,
 					},
+					{
+						$set: {
+							consumed: true,
+							consumedAt: Date.now(),
+							consumedBy: userID,
+						},
+					}
+				);
+
+				if (!inviteCodeDoc) {
+					logger.info(`Invalid invite code given: ${req.body.inviteCode}.`);
+					return res.status(401).json({
+						success: false,
+						description: `This invite code is not valid.`,
+					});
 				}
-			);
 
-			if (!inviteCodeDoc) {
-				logger.info(`Invalid invite code given: ${req.body.inviteCode}.`);
-				return res.status(401).json({
-					success: false,
-					description: `This invite code is not valid.`,
-				});
+				logger.info(`Consumed invite ${inviteCodeDoc.code}.`);
 			}
 
-			logger.info(`Consumed invite ${inviteCodeDoc.code}.`);
-
 			// if we get to this point, We're good to create the user.
 
 			const { newUser, newSettings } = await AddNewUser(
@@ -276,12 +279,12 @@ router.post(
 				body: user,
 			});
 		} catch (err) {
-			logger.error(
-				`Bailed on user creation ${req.body.username} with invite code ${req.body.inviteCode}.`,
-				{ err }
-			);
+			logger.error(`Bailed on user creation ${req.body.username}.`, { err });
+
+			if (!ServerConfig.ENABLE_PUBLIC_SIGNUP) {
+				await ReinstateInvite(req.body.inviteCode);
+			}
 
-			await ReinstateInvite(req.body.inviteCode);
 			await DecrementCounterValue("users");
 
 			return res.status(500).json({