input_wrapper.rb in upr-0.2.0

- old
+ new

@@ -7,14 +7,16 @@
   # that may wrap env["rack.input"] in the middleware chain.
   class InputWrapper < Struct.new(:app, :path_info, :frequency, :backend,
                                   :input, :pos, :seen, :content_length,
                                   :upload_id, :mtime)
 
+    include Params
+
     def initialize(app, options = {})
       super(app,
             Array(options[:path_info] || nil),
-            options[:frequency] || 3,
+            options[:frequency] || 1,
             options[:backend])
 
       # support :drb for compatibility with mongrel_upload_progress
       if options[:drb]
         backend and raise ArgumentError, ":backend and :drb are incompatible"
@@ -34,27 +36,22 @@
       if path_info.empty? || path_info.include?(env["PATH_INFO"])
         # benefit curl users...
         /\A100-continue\z/i =~ env['HTTP_EXPECT'] and return [ 100, {}, [] ]
 
         length = env["CONTENT_LENGTH"] and length = length.to_i
-        env["TRANSFER_ENCODING"] =~ %r{\Achunked\z}i and length = nil
-        if length.nil? || length > 0
-          req = Rack::Request.new(env)
-
-          # can't blindly parse params here since we don't want to read
-          # the POST body if there is one, so only parse stuff in the
-          # query string...
-          if uid = req.GET["upload_id"]
+        chunked = env["TRANSFER_ENCODING"] =~ %r{\Achunked\z}i and length = nil
+        if chunked || (length && length > 0)
+          if uid = extract_upload_id(env)
             return dup._call(env, uid, length)
           end
         end
       end
       app.call(env)
     end
 
     def _call(env, uid, length)
-      self.upload_id = env["upr.upload_id"] = uid
+      self.upload_id = uid
       self.mtime = self.pos = self.seen = 0
       self.input = env["rack.input"]
       env["rack.input"] = self
       self.content_length = length
       backend.start(upload_id, length)
@@ -62,24 +59,30 @@
       app.call(env)
     end
 
     def _incr(nr)
       self.pos += nr
+      _finish if content_length && pos >= content_length
       if (nr = pos - seen) > 0 && mtime <= (Time.now.to_i - frequency)
         backend.incr(upload_id, nr)
         self.seen = pos
         self.mtime = Time.now.to_i
       end
     end
 
+    def _finish
+      self.seen = backend.finish(upload_id).seen
+      self.content_length ||= self.seen
+    end
+
     def size
       rv = input.size
 
       # we had an unknown length and just had to read in everything to get it
       if content_length.nil?
         _incr(rv - seen)
-        self.content_length = rv
+        _finish
       end
       rv
     end
 
     def rewind
@@ -87,24 +90,25 @@
       input.rewind
     end
 
     def gets
       rv = input.gets
-      _incr(rv.size) unless rv.nil?
+      rv.nil? ? _finish : _incr(rv.size)
       rv
     end
 
     def read(*args)
       rv = input.read(*args)
-      _incr(rv.size) unless rv.nil?
+      rv.nil? || rv.size == 0 ? _finish : _incr(rv.size)
       rv
     end
 
     def each(&block)
       input.each do |chunk| # usually just a line
         _incr(chunk.size)
         yield chunk
       end
+      _finish
     end
 
   end
 end