1 files changed, 110 insertions, 0 deletions
diff --git a/lib/upr/input_wrapper.rb b/lib/upr/input_wrapper.rb
new file mode 100644
index 0000000..d267926
--- /dev/null
+++ b/lib/upr/input_wrapper.rb
@@ -0,0 +1,110 @@
+# -*- encoding: binary -*-
+
+module Upr
+
+  # the underlying middlware for for wrapping env["rack.input"],
+  # this should typically be installed before other middlewares
+  # 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)
+
+    def initialize(app, options = {})
+      super(app,
+            Array(options[:path_info] || nil),
+            options[:frequency] || 3,
+            options[:backend])
+
+      # support :drb for compatibility with mongrel_upload_progress
+      if options[:drb]
+        backend and raise ArgumentError, ":backend and :drb are incompatible"
+        require 'drb'
+        DRb.start_service
+        self.backend = DRbObject.new(nil, options[:drb])
+      elsif String === backend
+        # allow people to use strings in case their backend gets
+        # lazy-loaded (like an ActiveRecord model)
+        self.backend = eval(backend)
+      else
+        self.backend ||= Upr::Monitor.new
+      end
+    end
+
+    def call(env)
+      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"]
+            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.mtime = self.pos = self.seen = 0
+      self.input = env["rack.input"]
+      env["rack.input"] = self
+      self.content_length = length
+      backend.start(upload_id, length)
+
+      app.call(env)
+    end
+
+    def _incr(nr)
+      self.pos += nr
+      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 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
+      end
+      rv
+    end
+
+    def rewind
+      self.pos = 0
+      input.rewind
+    end
+
+    def gets
+      rv = input.gets
+      _incr(rv.size) unless rv.nil?
+      rv
+    end
+
+    def read(*args)
+      rv = input.read(*args)
+      _incr(rv.size) unless rv.nil?
+      rv
+    end
+
+    def each(&block)
+      input.each do |chunk| # usually just a line
+        _incr(chunk.size)
+        yield chunk
+      end
+    end
+
+  end
+end