Rewrote resource proxy code.

ckanext/resourceproxy/controller.py

@@ -1,7 +1,7 @@
-import urllib2
-import shutil
 from logging import getLogger
 
+import requests
+
 import ckan.logic as logic
 import ckan.lib.base as base
 
@@ -13,22 +13,49 @@ def proxy_resource(context, data_dict):
         log.info('Proxify resource {id}'.format(id=resource_id))
         resource = logic.get_action('resource_show')(context, {'id': resource_id})
         url = resource['url']
-        had_http_error = False
+
+        MAX_FILE_SIZE = 1024 * 1024 * 2  # 2MB
+        CHUNK_SIZE = 256
+
         try:
-            res = urllib2.urlopen(url)
-        except urllib2.HTTPError, error:
-            res = error
-            had_http_error = True
-        except urllib2.URLError, error:
-            details = "Could not proxy resource. %s" % str(error.reason)
-            base.abort(500, detail=details)
-        base.response.headers = res.headers
+            r = requests.get(url, prefetch=False)
+            r.raise_for_status()
+
+            # write body
+            cl = r.headers['content-length']
+            if cl and int(cl) > MAX_FILE_SIZE:
+                base.abort(500, '''Content is too large to be proxied.
+                    Allowed file size: {allowed}.
+                    Content-Length: {actual}'''.format(
+                        allowed=MAX_FILE_SIZE, actual=cl))
 
-        shutil.copyfileobj(res, base.response)
+            # write headers
+            base.response.headers = r.headers
 
-        # todo only change the status code, not the whole content
-        if had_http_error and hasattr(res, 'code'):
-            base.abort(res.code)
+            # we have to pretend that the response is not gzipped or deflated
+            # because we don't want request to unzip the content. There is a
+            # pull request pending that adds an option to iter_content.
+            r.headers['content-encoding'] = ''
+
+            length = 0
+            for chunk in r.iter_content(chunk_size=CHUNK_SIZE, decode_unicode=False):
+                base.response.body_file.write(chunk)
+                length += len(chunk)
+
+                if length >= MAX_FILE_SIZE:
+                    base.abort(500, headers={'content-encoding': ''},
+                        detail='Content is too large to be proxied.')
+
+        except requests.exceptions.HTTPError, error:
+            details = 'Could not proxy resource. %s' % str(error.response.reason)
+            base.abort(error.response.status_code, detail=details)
+        except requests.exceptions.ConnectionError, error:
+            details = '''Could not proxy resource because a
+                                connection error occurred. %s''' % str(error)
+            base.abort(500, detail=details)
+        except requests.exceptions.Timeout, error:
+            details = 'Could not proxy resource because the connection timed out.'
+            base.abort(500, detail=details)
 
 
 class ProxyController(base.BaseController):