From 6fe57d5e7896e6cfc865310596aa605758d0d3e0 Mon Sep 17 00:00:00 2001
From: Serdar Dogruyol <dogruyolserdar@gmail.com>
Date: Sat, 11 Feb 2017 15:33:42 +0200
Subject: [PATCH] Better file upload

---
 src/kemal/file_upload.cr     | 12 ++++++++++++
 src/kemal/helpers/helpers.cr | 30 ------------------------------
 src/kemal/param_parser.cr    | 32 ++++++++++++++++++++++++++++++--
 src/kemal/route_handler.cr   |  8 ++++++++
 4 files changed, 50 insertions(+), 32 deletions(-)
 create mode 100644 src/kemal/file_upload.cr

diff --git a/src/kemal/file_upload.cr b/src/kemal/file_upload.cr
new file mode 100644
index 00000000..4340a24e
--- /dev/null
+++ b/src/kemal/file_upload.cr
@@ -0,0 +1,12 @@
+# :nodoc:
+struct FileUpload
+  getter tmpfile : Tempfile
+  getter tmpfile_path : String
+  getter filename : String
+  getter meta : HTTP::FormData::FileMetadata
+  getter headers : HTTP::Headers
+
+  def initialize(@tmpfile, @tmpfile_path, @meta, @headers)
+    @filename = @meta.filename.not_nil!
+  end
+end
\ No newline at end of file
diff --git a/src/kemal/helpers/helpers.cr b/src/kemal/helpers/helpers.cr
index e7c00a67..9825cff2 100644
--- a/src/kemal/helpers/helpers.cr
+++ b/src/kemal/helpers/helpers.cr
@@ -72,33 +72,3 @@ end
 def gzip(status : Bool = false)
   add_handler HTTP::DeflateHandler.new if status
 end
-
-# :nodoc:
-struct UploadFile
-  getter field : String
-  getter data : IO::Delimited
-  getter meta : HTTP::FormData::FileMetadata
-  getter headers : HTTP::Headers
-
-  def initialize(@field, @data, @meta, @headers)
-  end
-end
-
-# Parses a multipart/form-data request. Yields an `UploadFile` object with `field`, `data`, `meta`, `headers` fields.
-# Consider the example below taking two image uploads as image1, image2. To get the relevant data
-# for each file you can use simple `if/switch` conditionals.
-#
-#   post "/upload" do |env|
-#     parse_multipart(env) do |f|
-#       image1 = f.data if f.field == "image1"
-#       image2 = f.data if f.field == "image2"
-#       puts f.meta
-#       puts f.headers
-#       "Upload complete"
-#     end
-#   end
-def parse_multipart(env)
-  HTTP::FormData.parse(env.request) do |field, data, meta, headers|
-    yield UploadFile.new field, data, meta, headers
-  end
-end
diff --git a/src/kemal/param_parser.cr b/src/kemal/param_parser.cr
index 5e28fede..dfb93e28 100644
--- a/src/kemal/param_parser.cr
+++ b/src/kemal/param_parser.cr
@@ -1,5 +1,6 @@
 require "json"
 require "uri"
+require "tempfile"
 
 module Kemal
   # ParamParser parses the request contents including query_params and body
@@ -8,14 +9,17 @@ module Kemal
   class ParamParser
     URL_ENCODED_FORM = "application/x-www-form-urlencoded"
     APPLICATION_JSON = "application/json"
+    MULTIPART_FORM   = "multipart/form-data"
     # :nodoc:
     alias AllParamTypes = Nil | String | Int64 | Float64 | Bool | Hash(String, JSON::Type) | Array(JSON::Type)
+    getter files
 
     def initialize(@request : HTTP::Request)
       @url = {} of String => String
       @query = HTTP::Params.new({} of String => Array(String))
       @body = HTTP::Params.new({} of String => Array(String))
       @json = {} of String => AllParamTypes
+      @files = {} of String => FileUpload
       @url_parsed = false
       @query_parsed = false
       @body_parsed = false
@@ -41,8 +45,16 @@ module Kemal
     {% end %}
 
     def parse_body
-      return if (@request.headers["Content-Type"]? =~ /#{URL_ENCODED_FORM}/).nil?
-      @body = parse_part(@request.body)
+      content_type = @request.headers["Content-Type"]?
+      return unless content_type
+      if content_type.try(&.starts_with?(URL_ENCODED_FORM))
+        @body = parse_part(@request.body)
+        return
+      end
+      if content_type.try(&.starts_with?(MULTIPART_FORM))
+        parse_file_upload
+        return
+      end
     end
 
     def parse_query
@@ -57,6 +69,22 @@ module Kemal
       end
     end
 
+    def parse_file_upload
+      HTTP::FormData.parse(@request) do |field, data, meta, headers|
+        next unless meta
+        filename = meta.filename
+        if !filename.nil?
+          tempfile = Tempfile.new(filename)
+          ::File.open(tempfile.path, "w") do |file|
+            IO.copy(data, file)
+          end
+          @files[field] = FileUpload.new(tmpfile: tempfile, tmpfile_path: tempfile.path, meta: meta, headers: headers)
+        else
+          @body[field] = data.gets_to_end
+        end
+      end
+    end
+
     # Parses JSON request body if Content-Type is `application/json`.
     # If request body is a JSON Hash then all the params are parsed and added into `params`.
     # If request body is a JSON Array it's added into `params` as `_json` and can be accessed
diff --git a/src/kemal/route_handler.cr b/src/kemal/route_handler.cr
index a04da21e..9031b419 100644
--- a/src/kemal/route_handler.cr
+++ b/src/kemal/route_handler.cr
@@ -34,6 +34,8 @@ module Kemal
       raise Kemal::Exceptions::RouteNotFound.new(context) unless context.route_defined?
       route = context.route_lookup.payload.as(Route)
       content = route.handler.call(context)
+    ensure
+      remove_tmpfiles(context)
       if Kemal.config.error_handlers.has_key?(context.response.status_code)
         raise Kemal::Exceptions::CustomException.new(context)
       end
@@ -41,6 +43,12 @@ module Kemal
       context
     end
 
+    private def remove_tmpfiles(context)
+      context.params.files.each do |field, file|
+        File.delete(file.tmpfile_path) if ::File.exists?(file.tmpfile_path)
+      end
+    end
+
     private def radix_path(method : String, path)
       "/#{method.downcase}#{path}"
     end