lib/symmetric_encryption/reader.rb in symmetric-encryption-0.6.0 vs lib/symmetric_encryption/reader.rb in symmetric-encryption-0.6.1

@@ -65,12 +65,11 @@
     #
     # # Example: Reading from a CSV file
     #
     # require 'fastercsv'
     # begin
-    #   # Must supply :row_sep for FasterCSV otherwise it will attempt to rewind the file etc.
-    #   csv = FasterCSV.new(SymmetricEncryption::Reader.open('csv_encrypted'), :row_sep => "\n")
+    #   csv = FasterCSV.new(SymmetricEncryption::Reader.open('csv_encrypted'))
     #   csv.each {|row| p row}
     # ensure
     #   csv.close if csv
     # end
     def self.open(filename_or_stream, options={}, &block)
@@ -90,33 +89,12 @@
 
     # Decrypt data before reading from the supplied stream
     def initialize(ios,options={})
       @ios         = ios
       @buffer_size = options.fetch(:buffer_size, 4096).to_i
-      @compressed  = nil
-      @read_buffer = ''
-
-      # Read first block and check for the header
-      buf = @ios.read(@buffer_size)
-      if buf.start_with?(SymmetricEncryption::MAGIC_HEADER)
-        # Header includes magic header and version byte
-        # Remove header and extract flags
-        header, flags = buf.slice!(0..MAGIC_HEADER_SIZE+1).unpack(MAGIC_HEADER_UNPACK)
-        @compressed = (flags & 0b1000_0000_0000_0000) != 0
-        @version = @compressed ? flags - 0b1000_0000_0000_0000 : flags
-      else
-        @version = options[:version]
-      end
-
-      # Use primary cipher by default, but allow a secondary cipher to be selected for encryption
-      @cipher = SymmetricEncryption.cipher(@version)
-      raise "Cipher with version:#{@version} not found in any of the configured SymmetricEncryption ciphers" unless @cipher
-      @stream_cipher = @cipher.send(:openssl_cipher, :decrypt)
-
-      # First call to #update should return an empty string anyway
-      @read_buffer << @stream_cipher.update(buf)
-      @read_buffer << @stream_cipher.final if @ios.eof?
+      @version = options[:version]
+      read_header
     end
 
     # Returns whether the stream being read is compressed
     #
     # Should be called before any reads are performed to determine if the file or
@@ -184,10 +162,11 @@
           buf = @ios.read || ''
           data << @stream_cipher.update(buf) if buf && buf.length > 0
           data << @stream_cipher.final
         end
       end
+      @pos += data.length
       data
     end
 
     # Reads a single decrypted line from the file up to and including the optional sep_string.
     # Raises EOFError on eof
@@ -207,10 +186,11 @@
       while (index = @read_buffer.index(sep_string)).nil? && !@ios.eof?
         read_block
       end
       index ||= -1
       data = @read_buffer.slice!(0..index)
+      @pos += data.length
       return nil if data.length == 0 && eof?
       data
     end
 
     # ios.each(sep_string="\n") {|line| block } => ios
@@ -229,23 +209,47 @@
     # Returns whether the end of file has been reached for this stream
     def eof?
       (@read_buffer.size == 0) && @ios.eof?
     end
 
-    # Return the approximate offset in bytes of the current input stream
-    # Since the encrypted data size does not match the unencrypted size
-    # this value cannot be guaranteed. Especially if compression is turned on
+    # Return the number of bytes read so far from the input stream
     def pos
-      @ios.pos - @read_buffer.size
+      @pos
     end
 
     # Rewind back to the beginning of the file
     def rewind
       @read_buffer = ''
       @ios.rewind
+      read_header
     end
 
     private
+
+    # Read the header from the file if present
+    def read_header
+      @compressed  = nil
+      @pos = 0
+
+      # Read first block and check for the header
+      buf = @ios.read(@buffer_size)
+      if buf.start_with?(MAGIC_HEADER)
+        # Header includes magic header and version byte
+        # Remove header and extract flags
+        header, flags = buf.slice!(0..MAGIC_HEADER_SIZE+1).unpack(MAGIC_HEADER_UNPACK)
+        @compressed = (flags & 0b1000_0000_0000_0000) != 0
+        @version = @compressed ? flags - 0b1000_0000_0000_0000 : flags
+      end
+
+      # Use primary cipher by default, but allow a secondary cipher to be selected for encryption
+      @cipher = SymmetricEncryption.cipher(@version)
+      raise "Cipher with version:#{@version.inspect} not found in any of the configured SymmetricEncryption ciphers" unless @cipher
+      @stream_cipher = @cipher.send(:openssl_cipher, :decrypt)
+
+      # First call to #update should return an empty string anyway
+      @read_buffer = @stream_cipher.update(buf)
+      @read_buffer << @stream_cipher.final if @ios.eof?
+    end
 
     # Read a block of data and append the decrypted data in the read buffer
     def read_block
       buf = @ios.read(@buffer_size)
       @read_buffer << @stream_cipher.update(buf) if buf && buf.length > 0