agent.rb in mechanize-2.3

- old
+ new

@@ -96,11 +96,11 @@
 
   # The cookies for this agent
   attr_accessor :cookie_jar
 
   # Responses larger than this will be written to a Tempfile instead of stored
-  # in memory.
+  # in memory.  Setting this to nil disables creation of Tempfiles.
   attr_accessor :max_file_buffer
 
   # :section: Utility
 
   # The context parses responses into pages
@@ -124,11 +124,11 @@
     @follow_meta_refresh      = false
     @follow_meta_refresh_self = false
     @gzip_enabled             = true
     @history                  = Mechanize::History.new
     @keep_alive               = true
-    @max_file_buffer          = 10240
+    @max_file_buffer          = 100_000 # 5MB for response bodies
     @open_timeout             = nil
     @post_connect_hooks       = []
     @pre_connect_hooks        = []
     @read_timeout             = nil
     @redirect_ok              = true
@@ -373,18 +373,11 @@
 
   def content_encoding_gunzip body_io
     log.debug('gzip response') if log
 
     zio = Zlib::GzipReader.new body_io
-    out_io = Tempfile.new 'mechanize-decode'
-    out_io.unlink
-    out_io.binmode
-
-    until zio.eof? do
-      out_io.write zio.read 16384
-    end
-
+    out_io = auto_io 'mechanize-gunzip', 16384, zio
     zio.finish
 
     return out_io
   rescue Zlib::Error
     log.error('unable to gunzip response, trying raw inflate') if log
@@ -398,10 +391,11 @@
       log.error("unable to gunzip response: #{e}") if log
       raise
     end
   ensure
     zio.close if zio and not zio.closed?
+    body_io.close unless body_io.closed?
   end
 
   ##
   # Decodes a deflate-encoded +body_io+.  If it cannot be decoded, raw inflate
   # is tried followed by raising an error.
@@ -419,10 +413,12 @@
       return inflate body_io, -Zlib::MAX_WBITS
     rescue Zlib::Error => e
       log.error("unable to inflate response: #{e}") if log
       raise
     end
+  ensure
+    body_io.close
   end
 
   def disable_keep_alive request
     request['connection'] = 'close' unless @keep_alive
   end
@@ -737,11 +733,14 @@
     message = "error handling content-encoding #{response['Content-Encoding']}:"
     message << " #{e.message} (#{e.class})"
     raise Mechanize::Error, message
   ensure
     begin
-      body_io.close! if Tempfile === body_io and out_io.path != body_io.path
+      if Tempfile === body_io and
+         (StringIO === out_io or out_io.path != body_io.path) then
+        body_io.close! 
+      end
     rescue IOError
       # HACK ruby 1.8 raises IOError when closing the stream
     end
   end
 
@@ -781,11 +780,11 @@
       redirects + 1 > @redirection_limit
 
     sleep delay
     @history.push(page, page.uri)
     fetch new_url, :get, {}, [],
-          Mechanize::Page.new(nil, {'content-type'=>'text/html'}), redirects
+          Mechanize::Page.new, redirects
   end
 
   def response_log response
     return unless log
 
@@ -802,11 +801,11 @@
   end
 
   def response_read response, request, uri
     content_length = response.content_length
 
-    if content_length and content_length > @max_file_buffer then
+    if use_tempfile? content_length then
       body_io = Tempfile.new 'mechanize-raw'
       body_io.unlink
       body_io.binmode if defined? body_io.binmode
     else
       body_io = StringIO.new
@@ -817,11 +816,11 @@
 
     begin
       response.read_body { |part|
         total += part.length
 
-        if StringIO === body_io and total > @max_file_buffer then
+        if StringIO === body_io and use_tempfile? total then
           new_io = Tempfile.new 'mechanize-raw'
           new_io.unlink
           new_io.binmode
 
           new_io.write body_io.string
@@ -1032,21 +1031,56 @@
     @http.idle_timeout = timeout
   end
 
   # :section: Utility
 
+  ##
+  # Creates a new output IO by reading +input_io+ in +read_size+ chunks.  If
+  # the output is over the max_file_buffer size a Tempfile with +name+ is
+  # created.
+  #
+  # If a block is provided, each chunk of +input_io+ is yielded for further
+  # processing.
+
+  def auto_io name, read_size, input_io
+    out_io = StringIO.new
+
+    out_io.set_encoding Encoding::BINARY if out_io.respond_to? :set_encoding
+
+    until input_io.eof? do
+      if StringIO === out_io and use_tempfile? out_io.size then
+        new_io = Tempfile.new name
+        new_io.unlink
+        new_io.binmode
+
+        new_io.write out_io.string
+        out_io = new_io
+      end
+
+      chunk = input_io.read read_size
+      chunk = yield chunk if block_given?
+
+      out_io.write chunk
+    end
+
+    out_io.rewind
+
+    out_io
+  end
+
   def inflate compressed, window_bits = nil
     inflate = Zlib::Inflate.new window_bits
-    out_io = Tempfile.new 'mechanize-decode'
 
-    until compressed.eof? do
-      out_io.write inflate.inflate compressed.read 1024
+    out_io = auto_io 'mechanize-inflate', 1024, compressed do |chunk|
+      inflate.inflate chunk
     end
 
     out_io.write inflate.finish
 
     out_io
+  ensure
+    inflate.close
   end
 
   def log
     @context.log
   end
@@ -1079,9 +1113,16 @@
     proxy_uri.port = port
     proxy_uri.user     = user if user
     proxy_uri.password = pass if pass
 
     @http.proxy = proxy_uri
+  end
+
+  def use_tempfile? size
+    return false unless @max_file_buffer
+    return false unless size
+
+    size >= @max_file_buffer
   end
 
 end
 