bigfile: better factor this for use with IO.copy_stream

syswrloop and it's filter arg is going away, but we're
stuck supporting bigfiles in read-only mode for eternity.

2 files changed, 74 insertions, 37 deletions

diff --git a/lib/mogilefs/bigfile.rb b/lib/mogilefs/bigfile.rb
index 34af8ee..1be36a5 100644
--- a/lib/mogilefs/bigfile.rb
+++ b/lib/mogilefs/bigfile.rb
@@ -1,12 +1,14 @@
 # -*- encoding: binary -*-
-require 'zlib'
-require 'digest/md5'
 require 'uri'
 require 'mogilefs/util'
+require 'mogilefs/network'
+
+# Used for reading deprecated "bigfile" objects generated by the deprecated
+# mogtool(1) utility.  Not recommended for new projects.
 
 module MogileFS::Bigfile
-  GZIP_HEADER = "\x1f\x8b".freeze # mogtool(1) has this
   # VALID_TYPES = %w(file tarball partition).map { |x| x.freeze }.freeze
+  include MogileFS::Network
 
   # returns a big_info hash if successful
   def bigfile_stat(key)
@@ -14,37 +16,17 @@ module MogileFS::Bigfile
   end
 
   # returns total bytes written and the big_info hash if successful, raises an
-  # exception if not wr_io is expected to be an IO-like object capable of
-  # receiving the syswrite method.
+  # exception if not.  wr_io is expected to be an IO-like object capable of
+  # receiving the write method.
   def bigfile_write(key, wr_io, opts = { :verify => false })
     info = bigfile_stat(key)
-    zi = nil
-    md5 = opts[:verify] ? Digest::MD5.new : nil
     total = 0
 
     # we only decode raw zlib deflated streams that mogtool (unfortunately)
     # generates.  tarballs and gzip(1) are up to to the application to decrypt.
-    filter = Proc.new do |buf|
-      if zi == nil
-        if info[:compressed] && info[:type] == 'file' &&
-             buf.length >= 2 && buf[0,2] != GZIP_HEADER
-          zi = Zlib::Inflate.new
-
-          # mogtool(1) seems to have a bug that causes it to generate bogus
-          # MD5s if zlib deflate is used.  Don't trust those MD5s for now...
-          md5 = nil
-        else
-          zi = false
-        end
-      end
-      buf ||= ''
-      if zi
-        zi.inflate(buf)
-      else
-        md5 << buf
-        buf
-      end
-    end if (info[:compressed] || md5)
+    if info[:compressed] || opts[:verify]
+      wr_io = MogileFS::Bigfile::Filter.new(wr_io, info, opts)
+    end
 
     info[:parts].each_with_index do |part,part_nr|
       next if part_nr == 0 # info[:parts][0] is always empty
@@ -58,16 +40,13 @@ module MogileFS::Bigfile
       end
 
       sock = http_read_sock(uris[0])
-      md5.reset if md5
-      w = sysrwloop(sock, wr_io, filter)
+      w = copy_stream(sock, wr_io)
 
-      if md5 && md5.hexdigest != part[:md5]
-        raise MogileFS::ChecksumMismatchError, "#{md5} != #{part[:md5]}"
-      end
+      wr_io.respond_to?(:md5_check!) and wr_io.md5_check!(part[:md5])
       total += w
     end
-
-    syswrite_full(wr_io, zi.finish) if zi
+    wr_io.flush
+    total += wr_io.flushed_bytes if wr_io.respond_to?(:flushed_bytes)
 
     [ total, info ]
   end
@@ -100,8 +79,8 @@ module MogileFS::Bigfile
 
       rv
     end
-
-end # module MogileFS::Bigfile
+end
+require "mogilefs/bigfile/filter"
 
 __END__
 # Copied from mogtool:
diff --git a/lib/mogilefs/bigfile/filter.rb b/lib/mogilefs/bigfile/filter.rb
new file mode 100644
index 0000000..e95a47e
--- /dev/null
+++ b/lib/mogilefs/bigfile/filter.rb
@@ -0,0 +1,58 @@
+# -*- encoding: binary -*-
+require 'zlib'
+require 'digest/md5'
+
+# Filter class to wrap IO objects and uncompress DEFLATE'd files
+#
+# This is used for reading "bigfile" objects generated by the
+# (deprecated) mogtool(1)
+class MogileFS::Bigfile::Filter
+  GZIP_HEADER = "\x1f\x8b"
+  INFLATABLE_TYPES = { "file" => true }
+  attr_reader :flushed_bytes
+
+  def initialize(io, info, opts)
+    @io = io
+    @info = info
+    @md5 = opts[:verify] ? Digest::MD5.new : nil
+    @zi = nil
+    @flushed_bytes = 0
+  end
+
+  def md5_check!(expect)
+    return unless @md5
+    current = @md5.hexdigest
+    current == expect or
+      raise MogileFS::ChecksumMismatchError, "#{current} != #{expect}"
+    @md5.reset
+  end
+
+  def flush
+    @flushed_bytes = @io.write(@zi.finish) if @zi
+    @io.flush
+  end
+
+  def write(buf)
+    if nil == @zi
+      if @info[:compressed] &&
+         INFLATABLE_TYPES.include?(@info[:type]) &&
+         buf.bytesize >= 2 &&
+         buf[0,2] != GZIP_HEADER
+
+        @zi = Zlib::Inflate.new
+
+        # mogtool(1) seems to have a bug that causes it to generate bogus
+        # MD5s if zlib deflate is used.  Don't trust those MD5s for now...
+        @md5 = nil
+      else
+        @zi = false
+      end
+    end
+    if @zi
+      buf = @zi.inflate(buf)
+    else
+      @md5 << buf
+    end
+    @io.write(buf)
+  end
+end