From ab4ae6fa1e2610637886f3a6107ec015ec19bbfd Mon Sep 17 00:00:00 2001
From: tycrek <josh.moore@jmoore.dev>
Date: Wed, 7 Jul 2021 13:17:55 -0600
Subject: [PATCH] added mimetype blocking (requires `npm run setup` after pull)

---
 MagicNumbers.json |  1 +
 README.md         |  3 ++-
 setup.js          |  7 +++++++
 storage.js        | 12 ++++++++++--
 4 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/MagicNumbers.json b/MagicNumbers.json
index 842c5e0..9ad1d5b 100644
--- a/MagicNumbers.json
+++ b/MagicNumbers.json
@@ -6,6 +6,7 @@
 	"CODE_UNAUTHORIZED": 401,
 	"CODE_NOT_FOUND": 404,
 	"CODE_PAYLOAD_TOO_LARGE": 413,
+	"CODE_UNSUPPORTED_MEDIA_TYPE": 415,
 	"CODE_INTERNAL_SERVER_ERROR": 500,
 	"KILOBYTES": 1024
 }
\ No newline at end of file
diff --git a/README.md b/README.md
index 3a60752..a0a85b9 100755
--- a/README.md
+++ b/README.md
@@ -23,7 +23,8 @@
 - ✔️ Upload images, videos, gifs, files
 - ✔️ Fancy embeds on Discord
 - ✔️ Seamless inline video embeds on Discord
-- ✔️ Upload log via customizable Discord Webhooks
+- ✔️ Personal upload log via customizable Discord Webhooks
+- ✔️ Mimetype blocking
 - ✔️ File deletion
 - ✔️ Usage metrics
 - ✔️ Thumbnail support
diff --git a/setup.js b/setup.js
index c0c50d8..c9b0455 100755
--- a/setup.js
+++ b/setup.js
@@ -12,6 +12,7 @@ const config = {
 	diskFilePath: 'uploads/',
 	saveWithDate: false,
 	saveAsOriginal: true,
+	mediaStrict: false,
 	s3enabled: false,
 	s3endpoint: 'sfo3.digitaloceanspaces.com',
 	s3bucket: 'bucket-name',
@@ -115,6 +116,12 @@ if (require.main === module) {
 				default: config.saveAsOriginal,
 				required: false
 			},
+			mediaStrict: {
+				description: 'Only allow uploads of media files (images, videos)',
+				type: 'boolean',
+				default: config.mediaStrict,
+				required: false
+			},
 			s3enabled: {
 				description: 'Enable uploading to S3 storage endpoints',
 				type: 'boolean',
diff --git a/storage.js b/storage.js
index 0d5032d..7886156 100644
--- a/storage.js
+++ b/storage.js
@@ -8,9 +8,11 @@ const Thumbnail = require('./thumbnails');
 const Vibrant = require('./vibrant');
 const Hash = require('./hash');
 const { getSafeExt, getDatedDirname, sanitize, generateId } = require('./utils');
-const { s3enabled, s3endpoint, s3bucket, s3accessKey, s3secretKey, saveAsOriginal, maxUploadSize } = require('./config.json');
+const { s3enabled, s3endpoint, s3bucket, s3accessKey, s3secretKey, saveAsOriginal, maxUploadSize, mediaStrict } = require('./config.json');
+const { CODE_UNSUPPORTED_MEDIA_TYPE } = require('./MagicNumbers.json');
 
 const ID_GEN_LENGTH = 32;
+const ALLOWED_MIMETYPES = /(image)|(video)\//;
 
 const s3 = new aws.S3({
 	endpoint: new aws.Endpoint(s3endpoint),
@@ -29,13 +31,19 @@ function getLocalFilename(req) {
 	return `${getDatedDirname()}/${saveAsOriginal ? req.file.originalname : req.file.sha1}`;
 }
 
-function processUploaded(req, _, next) {
+function processUploaded(req, res, next) {
 	// Fixes
 	req.file.mimetype = req.file.detectedMimeType;
 	req.file.originalname = sanitize(req.file.originalName);
 	req.file.randomId = generateId('random', ID_GEN_LENGTH, null, null);
 	req.file.deleteId = generateId('random', ID_GEN_LENGTH, null, null);
 
+	// Block the resource if the mimetype is not an image or video
+	if (mediaStrict && !ALLOWED_MIMETYPES.test(req.file.mimetype)) {
+		fs.remove(req.file.path.concat('.temp'));
+		return res.sendStatus(CODE_UNSUPPORTED_MEDIA_TYPE);
+	}
+
 	// Remove unwanted fields
 	delete req.file.fieldName;
 	delete req.file.originalName;