Allow source to be an IO object and remain backwards compatible to st…

…ring sources

.rspec

@@ -1,2 +1,2 @@
---order rand
 --color
+--format progress

README.md

@@ -10,6 +10,14 @@ hkdf.next_bytes(32)
  => "\f#\xF4b\x98\x9B\x7Fw>|/|k\xF4k\xB7\xB9\x11e\xC5\x92\xD1\fH\xFDG\x94vt\xB4\x14\xCE"
 ```
 
+You can also give an IO object as the source. It will be read in as a stream to generate the key. The optional arguement ```:read_size``` can be used to control how many bytes are read from the IO at a time.
+
+```ruby
+hkdf = HKDF.new(File.new('/tmp/filename'), :read_size => 512)
+hkdf.next_bytes(32)
+ => "\f#\xF4b\x98\x9B\x7Fw>|/|k\xF4k\xB7\xB9\x11e\xC5\x92\xD1\fH\xFDG\x94vt\xB4\x14\xCE"
+```
+
 The default algorithm is HMAC-SHA256, you can override this and other defaults by providing an options hash during construction.
 
 ```ruby

lib/hkdf.rb

@@ -1,16 +1,19 @@
 require 'openssl'
+require 'stringio'
 
 class HKDF
   def initialize(source, options = {})
-    options = {:algorithm => 'SHA256', :info => '', :salt => nil}.merge(options)
+    options = {:algorithm => 'SHA256', :info => '', :salt => nil, :read_size => nil}.merge(options)
+    source = StringIO.new(source) if source.is_a?(String)
 
     @digest = OpenSSL::Digest.new(options[:algorithm])
     @info = options[:info]
 
     salt = options[:salt]
     salt = 0.chr * @digest.digest_length if salt.nil? or salt.empty?
+    read_size = options[:read_size] || @digest.block_length
 
-    @prk = OpenSSL::HMAC.digest(@digest, salt, source)
+    @prk = _generate_prk(salt, source, read_size)
     @position = 0
     @blocks = []
     @blocks << ''
@@ -50,6 +53,14 @@ def next_hex_bytes(length)
     next_bytes(length).unpack('H*').first
   end
 
+  def _generate_prk(salt, source, read_size)
+    hmac = OpenSSL::HMAC.new(salt, @digest)
+    while block = source.read(read_size)
+      hmac.update(block)
+    end
+    hmac.digest
+  end
+
   def _generate_blocks(length)
     start = @blocks.size
     block_count = (length.to_f / @digest.digest_length).ceil

spec/hkdf_spec.rb

@@ -8,6 +8,26 @@
   end
 
   describe 'initialize' do
+    it 'accepts an IO or a string as a source' do
+      output1 = HKDF.new(@source).next_bytes(32)
+      output2 = HKDF.new(StringIO.new(@source)).next_bytes(32)
+      output1.should == output2
+    end
+
+    it 'reads in an IO at a given read size' do
+      io = StringIO.new(@source)
+      io.should_receive(:read).with(1)
+
+      HKDF.new(io, :read_size => 1)
+    end
+
+    it 'reads in the whole IO' do
+      hkdf1 = HKDF.new(@source, :read_size => 1)
+      hkdf2 = HKDF.new(@source)
+
+      hkdf1.next_bytes(32).should == hkdf2.next_bytes(32)
+    end
+
     it 'defaults the algorithm to SHA-256' do
       HKDF.new(@source).algorithm.should == 'SHA256'
     end
@@ -46,15 +66,15 @@
 
   describe 'next_bytes' do
     it 'raises an error if requested size is > max_length' do
-      expect { @hkdf.next_bytes(@hkdf.max_length + 1) }.should raise_error(RangeError, /requested \d+ bytes, only \d+ available/)
-      expect { @hkdf.next_bytes(@hkdf.max_length) }.should_not raise_error(RangeError)
+      expect { @hkdf.next_bytes(@hkdf.max_length + 1) }.to raise_error(RangeError, /requested \d+ bytes, only \d+ available/)
+      expect { @hkdf.next_bytes(@hkdf.max_length) }.to_not raise_error(RangeError)
     end
 
     it 'raises an error if requested size + current position is > max_length' do
       expect do
         @hkdf.next_bytes(32)
         @hkdf.next_bytes(@hkdf.max_length - 31)
-      end.should raise_error(RangeError, /requested \d+ bytes, only \d+ available/)
+      end.to raise_error(RangeError, /requested \d+ bytes, only \d+ available/)
     end
 
     it 'advances the stream position' do
@@ -86,8 +106,8 @@
     end
 
     it 'raises an error if requested to seek past end of stream' do
-      expect { @hkdf.seek(@hkdf.max_length + 1) }.should raise_error(RangeError, /cannot seek past \d+/)
-      expect { @hkdf.seek(@hkdf.max_length) }.should_not raise_error(RangeError)
+      expect { @hkdf.seek(@hkdf.max_length + 1) }.to raise_error(RangeError, /cannot seek past \d+/)
+      expect { @hkdf.seek(@hkdf.max_length) }.to_not raise_error(RangeError)
     end
   end
 

spec/spec_helper.rb

@@ -1,5 +1,10 @@
 require 'hkdf'
 
+RSpec.configure do |config|
+  config.treat_symbols_as_metadata_keys_with_true_values = true
+  config.order = 'random'
+end
+
 def test_vectors
   test_lines = File.readlines('spec/test_vectors.txt').map(&:strip).reject(&:empty?)