fix(api:train): validate mime type and result of database.get.fileByF…

…ilename (#123)
jakowenko · Oct 12, 2021 · d5e050f · d5e050f
1 parent d607d76
commit d5e050f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/api/src/controllers/train.controller.js b/api/src/controllers/train.controller.js
@@ -99,7 +99,11 @@ module.exports.add = async (req, res) => {
   if (req.files) {
     await Promise.all(
       req.files.map(async (obj) => {
-        const { originalname, buffer } = obj;
+        const { originalname, buffer, mimetype } = obj;
+        if (!['image/jpeg', 'image/png'].includes(mimetype)) {
+          console.warn(`training incorrect mime type: ${mimetype}`);
+          return;
+        }
         const ext = `.${originalname.split('.').pop()}`;
         const filename = `${originalname.replace(ext, '')}-${time.unix()}${ext}`;
         await filesystem.writer(`${STORAGE.PATH}/train/${name}/${filename}`, buffer);
@@ -123,7 +127,7 @@ module.exports.add = async (req, res) => {
     }));
   }
 
-  train.add(name, { files });
+  if (files.length) train.add(name, { files });
   res.send({ message: `training queued for ${name}` });
 };
 

diff --git a/api/src/util/train.util.js b/api/src/util/train.util.js
@@ -79,7 +79,7 @@ module.exports.add = async (name, opts = {}) => {
   const { ids, files } = opts;
   await database.resync.files();
   const queue = files
-    ? files.map((obj) => database.get.fileByFilename(obj.name, obj.filename))
+    ? files.map((obj) => database.get.fileByFilename(obj.name, obj.filename)).filter((obj) => obj)
     : ids
     ? database.get.filesById(ids)
     : database.get.untrained(name);