Merge pull request #1243 from reidmv/util-downloader

Implmement R10K::Util::Downloader

lib/r10k/util/downloader.rb

@@ -0,0 +1,134 @@
+require 'digest'
+require 'net/http'
+
+module R10K
+  module Util
+
+    # Utility mixin for classes that need to download files
+    module Downloader
+
+      # Downloader objects need to checksum downloaded or saved content. The
+      # algorithm used to perform this checksumming (and therefore the kinds of
+      # checksums returned by various methods) is reported by this method.
+      #
+      # @return [Symbol] The checksum algorithm the downloader uses
+      def checksum_algorithm
+        @checksum_algorithm ||= :SHA256
+      end
+
+      private
+
+      # Set the checksum algorithm the downloader should use. It should be a
+      # symbol, and a valid Ruby 'digest' library algorithm. The default is
+      # :SHA256.
+      #
+      # @param algorithm [Symbol] The checksum algorithm the downloader should use
+      def checksum_algorithm=(algorithm)
+        @checksum_algorithm = algorithm
+      end
+
+      CHUNK_SIZE = 64 * 1024 # 64 kb
+
+      # @param src_uri [URI] The URI to download from
+      # @param dst_file [String] The file or path to save to
+      # @return [String] The downloaded file's hex digest
+      def download(src_uri, dst_file)
+        digest = Digest(checksum_algorithm).new
+        http_get(src_uri) do |resp|
+          File.open(dst_file, 'wb') do |output_stream|
+            resp.read_body do |chunk|
+              output_stream.write(chunk)
+              digest.update(chunk)
+            end
+          end
+        end
+
+        digest.hexdigest
+      end
+
+      # @param src_file The file or path to copy from
+      # @param dst_file The file or path to copy to
+      # @return [String] The copied file's sha256 hex digest
+      def copy(src_file, dst_file)
+        digest = Digest(checksum_algorithm).new
+        File.open(src_file, 'rb') do |input_stream|
+          File.open(dst_file, 'wb') do |output_stream|
+            until input_stream.eof?
+              chunk = input_stream.read(CHUNK_SIZE)
+              output_stream.write(chunk)
+              digest.update(chunk)
+            end
+          end
+        end
+
+        digest.hexdigest
+      end
+
+      # Start a Net::HTTP::Get connection, then yield the Net::HTTPSuccess object
+      # to the caller's block. Follow redirects if Net::HTTPRedirection responses
+      # are encountered, and use a proxy if directed.
+      #
+      # @param uri [URI] The URI to download the file from
+      # @param redirect_limit [Integer] How many redirects to permit before failing
+      # @param proxy [URI, String] The URI to use as a proxy
+      def http_get(uri, redirect_limit: 10, proxy: nil, &block)
+        raise "HTTP redirect too deep" if redirect_limit.zero?
+
+        session = Net::HTTP.new(uri.host, uri.port, *proxy_to_array(proxy))
+        session.use_ssl = true if uri.scheme == 'https'
+        session.start
+
+        begin
+          session.request_get(uri) do |response|
+            case response
+            when Net::HTTPRedirection
+              redirect = response['location']
+              session.finish
+              return http_get(URI.parse(redirect), redirect_limit: redirect_limit - 1, proxy: proxy, &block)
+            when Net::HTTPSuccess
+              yield response
+            else
+              raise "Unexpected response code #{response.code}: #{response}"
+            end
+          end
+        ensure
+          session.finish if session.active?
+        end
+      end
+
+      # Helper method to translate a proxy URI to array arguments for
+      # Net::HTTP#new. A nil argument returns nil array elements.
+      def proxy_to_array(proxy_uri)
+        if proxy_uri
+          px = proxy_uri.is_a?(URI) ? proxy_uri : URI.parse(proxy_uri)
+          [px.host, px.port, px.user, px.password]
+        else
+          [nil, nil, nil, nil]
+        end
+      end
+
+      # Return the sha256 digest of the file at the given path
+      #
+      # @param path [String] The path to the file
+      # @return [String] The file's sha256 hex digest
+      def file_digest(path)
+        File.open(path) do |file|
+          reader_digest(file)
+        end
+      end
+
+      # Return the sha256 digest of the readable data
+      #
+      # @param reader [String] An object that responds to #read
+      # @return [String] The read data's sha256 hex digest
+      def reader_digest(reader)
+        digest = Digest(checksum_algorithm).new
+        while chunk = reader.read(CHUNK_SIZE)
+          digest.update(chunk)
+        end
+
+        digest.hexdigest
+      end
+    end
+  end
+end

spec/unit/util/downloader_spec.rb

@@ -0,0 +1,98 @@
+require 'spec_helper'
+require 'r10k/util/downloader'
+
+describe R10K::Util::Downloader do
+
+  subject(:downloader) do
+    subj = Object.new
+    subj.extend(R10K::Util::Downloader)
+    subj.singleton_class.class_eval { public :download }
+    subj.singleton_class.class_eval { public :http_get }
+    subj.singleton_class.class_eval { public :file_digest }
+    subj
+  end
+
+  let(:tmpdir) { Dir.mktmpdir }
+  after(:each) { FileUtils.remove_entry_secure(tmpdir) }
+
+  describe 'http_get' do
+    let(:src_url) { 'https://example.com' }
+    let(:dst_file) { File.join(tmpdir, 'test.out') }
+    let(:tarball_uri) { URI('http://tarball.example.com/tarball.tar.gz') }
+    let(:redirect_uri) { URI('http://redirect.example.com/redirect') }
+    let(:proxy_uri) { URI('http://user:password@proxy.example.com') }
+
+    it 'downloads a simple file' do
+      mock_session = instance_double('Net::HTTP', active?: true)
+      tarball_response = instance_double('Net::HTTPSuccess')
+
+      expect(Net::HTTP).to receive(:new).with(tarball_uri.host, any_args).and_return(mock_session)
+      expect(Net::HTTPSuccess).to receive(:===).with(tarball_response).and_return(true)
+
+      expect(mock_session).to receive(:request_get).and_yield(tarball_response)
+      expect(mock_session).to receive(:start).once
+      expect(mock_session).to receive(:finish).once
+
+      expect { |b| downloader.http_get(tarball_uri, &b) }.to yield_with_args(tarball_response)
+    end
+
+    it 'follows redirects' do
+      mock_session_1 = instance_double('Net::HTTP', active?: false)
+      mock_session_2 = instance_double('Net::HTTP', active?: true)
+      redirect_response = instance_double('Net::HTTPRedirection')
+      tarball_response = instance_double('Net::HTTPSuccess')
+
+      expect(Net::HTTP).to receive(:new).with(redirect_uri.host, any_args).and_return(mock_session_1).once
+      expect(Net::HTTP).to receive(:new).with(tarball_uri.host, any_args).and_return(mock_session_2).once
+      expect(Net::HTTPRedirection).to receive(:===).with(redirect_response).and_return(true)
+      expect(Net::HTTPSuccess).to receive(:===).with(tarball_response).and_return(true)
+      allow(Net::HTTPRedirection).to receive(:===).and_call_original
+
+      expect(mock_session_1).to receive(:request_get).and_yield(redirect_response)
+      expect(mock_session_2).to receive(:request_get).and_yield(tarball_response)
+
+      # The redirect response should be queried for the redirect location
+      expect(redirect_response).to receive(:[]).with('location').and_return(tarball_uri.to_s)
+
+      # Both sessions should start and finish cleanly
+      expect(mock_session_1).to receive(:start).once
+      expect(mock_session_1).to receive(:finish).once
+      expect(mock_session_2).to receive(:start).once
+      expect(mock_session_2).to receive(:finish).once
+
+      expect { |b| downloader.http_get(redirect_uri, &b) }.to yield_with_args(tarball_response)
+    end
+
+    it 'can use a proxy' do
+      mock_session = instance_double('Net::HTTP', active?: true)
+
+      expect(Net::HTTP).to receive(:new)
+                       .with(tarball_uri.host,
+                             tarball_uri.port,
+                             proxy_uri.host,
+                             proxy_uri.port,
+                             proxy_uri.user,
+                             proxy_uri.password,
+                             any_args)
+                       .and_return(mock_session)
+
+      expect(mock_session).to receive(:request_get).and_return(:not_yielded)
+      expect(mock_session).to receive(:start).once
+      expect(mock_session).to receive(:finish).once
+
+      downloader.http_get(tarball_uri, proxy: proxy_uri)
+    end
+  end
+
+  describe 'checksums' do
+    let(:fixture_checksum) { '0bcea17aa0c5e868c18f0fa042feda770e47c1a4223229f82116ccb3ca33c6e3' }
+    let(:fixture_tarball) do
+      File.expand_path('spec/fixtures/integration/git/puppet-boolean-bare.tar', PROJECT_ROOT)
+    end
+
+    it 'checksums files' do
+      expect(downloader.file_digest(fixture_tarball)).to eql(fixture_checksum)
+    end
+  end
+end
+