From 26e009b474c7b940e241d4d1ae0b33b51c28b823 Mon Sep 17 00:00:00 2001
From: Adam Schrader <adam.dot.schrader@gmail.com>
Date: Fri, 28 Apr 2017 13:53:58 -0400
Subject: [PATCH] Add "preserverExtension" option handler

Misc reformatting.
---
 README.md    |  1 +
 lib/index.js | 33 ++++++++++++++++++++++++++++-----
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 5b1a098..0f6e3a6 100644
--- a/README.md
+++ b/README.md
@@ -121,6 +121,7 @@ Pass in non-Busboy options directly to the middleware. These are express-fileupl
 Option | Acceptable&nbsp;Values | Details
 --- | --- | ---
 safeFileNames | <ul><li><code>false</code>&nbsp;**(default)**</li><li><code>true</code></li><li>regex</li></ul> | Strips characters from the upload's filename. You can use custom regex to determine what to strip. If set to `true`, non-alphanumeric characters _except_ dashes and underscores will be stripped. This option is off by default.<br /><br />**Example #1 (strip slashes from file names):** `app.use(fileUpload({ safeFileNames: /\\/g }))`<br />**Example #2:** `app.use(fileUpload({ safeFileNames: true }))`
+preserveExtension | <ul><li><code>false</code>&nbsp;**(default)**</li><li><code>true</code></li><li><code>*Number*</code></li></ul> | Preserves filename extension when using <code>safeFileNames</code> option. If set to <code>true</code>, will default to an extension length of 3. If set to <code>*Number*</code>, this will be the max allowable extension length. If an extension is smaller than the extension length, it remains untouched. If the extension is longer, it is shifted.<br /><br />**Example #1 (true):**<br /><code>app.use(fileUpload({ safeFileNames: true, preserveExtension: true }));</code><br />*myFileName.ext* --> *myFileName.ext*<br /><br />**Example #2 (max extension length 2, extension truncated):**<br /><code>app.use(fileUpload({ safeFileNames: true, preserveExtension: 2 }));</code><br />*myFileName.ext* --> *myFileNamee.xt*
 
 # Help Wanted
 Pull Requests are welcomed!
diff --git a/lib/index.js b/lib/index.js
index c34d443..885025f 100644
--- a/lib/index.js
+++ b/lib/index.js
@@ -15,7 +15,7 @@ module.exports = function(options) {
 
   return function(req, res, next) {
     if (!hasBody(req) || !hasAcceptableMethod(req) || !hasAcceptableMime(req))
-        return next();
+      return next();
 
     processMultipart(options, req, res, next);
   };
@@ -85,14 +85,37 @@ function processMultipart(options, req, res, next) {
       // see: https://github.com/richardgirges/express-fileupload/issues/14
       // firefox uploads empty file in case of cache miss when f5ing page.
       // resulting in unexpected behavior. if there is no file data, the file is invalid.
-      if(!buf.length)
+      if (!buf.length)
         return;
 
       if (options.safeFileNames) {
+        let extensionLength = 3;
+        let extension = '';
+
         if (typeof options.safeFileNames === 'object')
           safeFileNameRegex = options.safeFileNames;
 
-        filename = filename.replace(safeFileNameRegex, '');
+        if (options.preserveExtension) {
+          if (typeof options.preserveExtension === 'number')
+            extensionLength = options.preserveExtension;
+
+          let filenameParts = filename.split('.');
+          let filenamePartsLen = filenameParts.length;
+          if (filenamePartsLen > 1) {
+            extension = filenameParts.pop();
+
+            if (extension.length > extensionLength) {
+              filenameParts[filenameParts.length - 1] +=
+                '.' + extension.substr(0, extension.length - extensionLength);
+              extension = extension.substr(-extensionLength);
+            }
+
+            extension = '.' + extension.replace(safeFileNameRegex, '');
+            filename = filenameParts.join('.');
+          }
+        }
+
+        filename = filename.replace(safeFileNameRegex, '').concat(extension);
       }
 
       let newFile = {
@@ -123,9 +146,9 @@ function processMultipart(options, req, res, next) {
       } else {
         // Array fields
         if (req.files[fieldname] instanceof Array)
-            req.files[fieldname].push(newFile);
+          req.files[fieldname].push(newFile);
         else
-            req.files[fieldname] = [req.files[fieldname], newFile];
+          req.files[fieldname] = [req.files[fieldname], newFile];
       }
     });
   });