lib/link_thumbnailer/processor.rb in link_thumbnailer-3.3.2 vs lib/link_thumbnailer/processor.rb in link_thumbnailer-3.4.0

@@ -15,10 +15,16 @@
       @http   = ::Net::HTTP::Persistent.new
 
       super(config)
     end
 
+    def start(url)
+      result = call(url)
+      shutdown
+      result
+    end
+
     def call(url = '', redirect_count = 0, headers = {})
       self.url        = url
       @redirect_count = redirect_count
 
       raise ::LinkThumbnailer::RedirectLimit if too_many_redirections?
@@ -26,16 +32,20 @@
       with_valid_url do
         set_http_headers(headers)
         set_http_options
         perform_request
       end
-    rescue ::Net::HTTPExceptions, ::SocketError, ::Timeout::Error => e
+    rescue ::Net::HTTPExceptions, ::SocketError, ::Timeout::Error, ::Net::HTTP::Persistent::Error => e
       raise ::LinkThumbnailer::HTTPError.new(e.message)
     end
 
     private
 
+    def shutdown
+      http.shutdown
+    end
+
     def with_valid_url
       raise ::LinkThumbnailer::BadUriFormat unless valid_url_format?
       yield if block_given?
     end
 
@@ -51,11 +61,11 @@
       http.read_timeout = http_read_timeout
       http.proxy = :ENV
     end
 
     def perform_request
-      response          = http.request(url)
+      response          = request_in_chunks
       headers           = {}
       headers['Cookie'] = response['Set-Cookie'] if response['Set-Cookie'].present?
 
       raise ::LinkThumbnailer::FormatNotSupported.new(response['Content-Type']) unless valid_response_format?(response)
 
@@ -71,10 +81,23 @@
       else
         response.error!
       end
     end
 
+    def request_in_chunks
+      body     = String.new
+      response = http.request(url) do |resp|
+        raise ::LinkThumbnailer::DownloadSizeLimit if too_big_download_size?(resp.content_length)
+        resp.read_body do |chunk|
+          body.concat(chunk)
+          raise ::LinkThumbnailer::DownloadSizeLimit if too_big_download_size?(body.length)
+        end
+      end
+      response.body = body
+      response
+    end
+
     def resolve_relative_url(location)
       location.start_with?('http') ? location : build_absolute_url_for(location)
     end
 
     def build_absolute_url_for(relative_url)
@@ -99,10 +122,14 @@
 
     def ssl_required?
       config.verify_ssl
     end
 
+    def download_size_limit
+      config.download_size_limit
+    end
+
     def too_many_redirections?
       redirect_count > redirect_limit
     end
 
     def valid_url_format?
@@ -116,9 +143,13 @@
       return true if response['Content-Type'] =~ /application\/xhtml\+xml/
       return true if response['Content-Type'] =~ /application\/xml/
       return true if response['Content-Type'] =~ /text\/xml/
       return true if response['Content-Type'] =~ /text\/plain/
       false
+    end
+
+    def too_big_download_size?(size)
+      size.to_i > download_size_limit.to_i
     end
 
     def url=(url)
       @url = ::URI.parse(url.to_s)
     end