From aeb4450a5693743d791261804289b6060c067eff Mon Sep 17 00:00:00 2001
From: Isaac <git@eartharoid.me>
Date: Fri, 10 Mar 2023 23:46:24 +0000
Subject: [PATCH] feat: oauth2 callback redirect (closes #333)

---
 src/http.js                 | 15 ++++++++++++++-
 src/routes/auth/callback.js | 19 +++++++++----------
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/src/http.js b/src/http.js
index fb73b15..95ad08f 100644
--- a/src/http.js
+++ b/src/http.js
@@ -1,6 +1,6 @@
 const fastify = require('fastify')({ trustProxy: process.env.HTTP_TRUST_PROXY === 'true' });
 const oauth = require('@fastify/oauth2');
-const { domain } = require('./lib/http');
+const { randomBytes } = require('crypto');
 const { short } = require('leeks.js');
 const { join } = require('path');
 const { files } = require('node-dir');
@@ -17,8 +17,16 @@ module.exports = async client => {
 	});
 
 	// oauth2 plugin
+	fastify.states = new Map();
 	fastify.register(oauth, {
 		callbackUri: `${process.env.HTTP_EXTERNAL}/auth/callback`,
+		checkStateFunction: (state, callback) => {
+			if (fastify.states.has(state)) {
+				callback();
+				return;
+			}
+			callback(new Error('Invalid state'));
+		},
 		credentials: {
 			auth: oauth.DISCORD_CONFIGURATION,
 			client: {
@@ -26,6 +34,11 @@ module.exports = async client => {
 				secret: process.env.DISCORD_SECRET,
 			},
 		},
+		generateStateFunction: req => {
+			const state = randomBytes(12).toString('hex');
+			fastify.states.set(state, req.query.r);
+			return state;
+		},
 		name: 'discord',
 		scope: ['applications.commands.permissions.update', 'guilds', 'identify'],
 		startRedirectPath: '/auth/login',
diff --git a/src/routes/auth/callback.js b/src/routes/auth/callback.js
index c319b37..2d04e9e 100644
--- a/src/routes/auth/callback.js
+++ b/src/routes/auth/callback.js
@@ -3,15 +3,15 @@ const { domain } = require('../../lib/http');
 module.exports.get = () => ({
 	handler: async function (req, res) { // MUST NOT use arrow function syntax
 		const {
-			access_token,
-			expires_in,
+			access_token: accessToken,
+			expires_in: expiresIn,
 		} = await this.discord.getAccessTokenFromAuthorizationCodeFlow(req);
-		const user = await (await fetch('https://discordapp.com/api/users/@me', { headers: { 'Authorization': `Bearer ${access_token}` } })).json();
+		const user = await (await fetch('https://discordapp.com/api/users/@me', { headers: { 'Authorization': `Bearer ${accessToken}` } })).json();
 		const payload = {
-			access_token,
+			accessToken,
 			avatar: `https://cdn.discordapp.com/avatars/${user.id}/${user.avatar}.webp`,
 			discriminator: user.discriminator,
-			expiresAt: Date.now() + (expires_in * 1000),
+			expiresAt: Date.now() + (expiresIn * 1000),
 			id: user.id,
 			locale: user.locale,
 			username: user.username,
@@ -20,15 +20,14 @@ module.exports.get = () => ({
 		const token = this.jwt.sign({ payload });
 		res
 			.setCookie('token', token, {
-				domain: domain,
+				domain,
 				httpOnly: true,
-				maxAge: 604800, // seconds, not milliseconds
+				maxAge: expiresIn,
 				path: '/',
 				sameSite: true,
 				secure: false,
 			})
-			// .redirect('/settings')
-			.type('text/html')
-			.send('<a href="/settings">/settings</a>'); // temp fix: redirecting causes weird discord<->callback loop, probably caching?
+			.redirect(this.states.get(req.query.state) || '/');
+		this.states.delete(req.query.state);
 	},
 });
\ No newline at end of file