From 26574d05112ba49f9311fc93085e01c14073aaf5 Mon Sep 17 00:00:00 2001
From: Matteo Collina <hello@matteocollina.com>
Date: Thu, 5 Jul 2018 14:38:11 +0200
Subject: [PATCH] Wait for all files to be completed before calling the
 handler.

---
 index.js     | 37 +++++++++++++++++++++++++++++++--
 package.json |  1 +
 test.js      | 58 +++++++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 93 insertions(+), 3 deletions(-)

diff --git a/index.js b/index.js
index d6d96ab8..dbb900d2 100644
--- a/index.js
+++ b/index.js
@@ -3,6 +3,7 @@
 const fp = require('fastify-plugin')
 const Busboy = require('busboy')
 const kMultipart = Symbol('multipart')
+const eos = require('end-of-stream')
 
 function setMultipart (req, done) {
   // nothing to do, it will be done by the Request.multipart object
@@ -46,15 +47,27 @@ function fastifyMultipart (fastify, options, done) {
       busboyOptions[keys[i]] = options[keys[i]]
     }
     const stream = new Busboy(busboyOptions)
+    var completed = false
+    var files = 0
+    var count = 0
+    var callDoneOnNextEos = false
 
     req.on('error', function (err) {
       stream.destroy()
-      done(err)
+      if (!completed) {
+        completed = true
+        done(err)
+      }
     })
 
     stream.on('finish', function () {
       log.debug('finished multipart parsing')
-      done()
+      if (!completed && count === files) {
+        completed = true
+        setImmediate(done)
+      } else {
+        callDoneOnNextEos = true
+      }
     })
 
     stream.on('file', wrap)
@@ -63,9 +76,29 @@ function fastifyMultipart (fastify, options, done) {
 
     function wrap (field, file, filename, encoding, mimetype) {
       log.debug({ field, filename, encoding, mimetype }, 'parsing part')
+      files++
+      eos(file, waitForFiles)
       handler(field, file, filename, encoding, mimetype)
     }
 
+    function waitForFiles (err) {
+      if (err) {
+        completed = true
+        done(err)
+        return
+      }
+
+      if (completed) {
+        return
+      }
+
+      ++count
+      if (callDoneOnNextEos && count === files) {
+        completed = true
+        done()
+      }
+    }
+
     return stream
   }
 
diff --git a/package.json b/package.json
index 3bdb1773..4f27cde1 100644
--- a/package.json
+++ b/package.json
@@ -5,6 +5,7 @@
   "main": "index.js",
   "dependencies": {
     "busboy": "^0.2.14",
+    "end-of-stream": "^1.4.1",
     "fastify-plugin": "^1.0.0"
   },
   "devDependencies": {
diff --git a/test.js b/test.js
index b5f39c49..518129cf 100644
--- a/test.js
+++ b/test.js
@@ -1,5 +1,5 @@
 'use strict'
-
+const os = require('os')
 const test = require('tap').test
 const FormData = require('form-data')
 const Fastify = require('fastify')
@@ -80,6 +80,62 @@ test('should parse forms', function (t) {
   })
 })
 
+test('should call finished when both files are pumped', function (t) {
+  t.plan(8)
+
+  const fastify = Fastify()
+  t.tearDown(fastify.close.bind(fastify))
+
+  fastify.register(multipart)
+
+  fastify.post('/', function (req, reply) {
+    var fileCount = 0
+    t.ok(req.isMultipart())
+
+    req.multipart(handler, function (err) {
+      t.error(err)
+      t.equal(fileCount, 2)
+      reply.code(200).send()
+    })
+
+    function handler (field, file, filename, encoding, mimetype) {
+      const saveTo = path.join(os.tmpdir(), path.basename(filename))
+      pump(file, fs.createWriteStream(saveTo), function (err) {
+        t.error(err)
+        fileCount++
+      })
+    }
+  })
+
+  fastify.listen(0, function () {
+    // request
+    var form = new FormData()
+    var opts = {
+      protocol: 'http:',
+      hostname: 'localhost',
+      port: fastify.server.address().port,
+      path: '/',
+      headers: form.getHeaders(),
+      method: 'POST'
+    }
+
+    var req = http.request(opts, (res) => {
+      t.equal(res.statusCode, 200)
+      res.resume()
+      res.on('end', () => {
+        t.pass('res ended successfully')
+      })
+    })
+    form.append('upload', fs.createReadStream(filePath))
+    form.append('upload2', fs.createReadStream(filePath))
+    form.append('hello', 'world')
+    form.append('willbe', 'dropped')
+    pump(form, req, function (err) {
+      t.error(err, 'client pump: no err')
+    })
+  })
+})
+
 test('should error if it is not multipart', function (t) {
   t.plan(4)