Hard-cap image size to 256 kB

.gitignore

@@ -1 +1,2 @@
 env
+*.egg-info

setup.py

@@ -0,0 +1,3 @@
+from setuptools import setup, find_packages
+
+setup(name='gip', packages=find_packages())

tests/test_v1.py

@@ -0,0 +1,61 @@
+# -*- coding: utf-8 -*-
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+from cStringIO import StringIO
+from os.path import join
+from zipfile import ZipFile
+
+import pytest
+from aspen.testing.client import Client
+from gip import website
+
+
+@pytest.fixture
+def client():
+    client = Client(www_root='www', project_root='.')
+    client._website = website
+    yield client
+
+@pytest.fixture
+def image():
+    yield open(join('tests', 'image.jpg')).read()
+
+
+def zipfile(raw):
+    return ZipFile(StringIO(raw), mode='r')
+
+def filenames(zipfile):
+    return [n.filename for n in zipfile.filelist]
+
+@pytest.fixture
+def vary_length(client, image):
+    def vary_length(length, fail=True):
+        method = client.PxST if fail else client.POST
+        return method( '/v1'
+                     , body=image
+                     , content_type='image/jpeg'
+                     , HTTP_CONTENT_LENGTH=str(length)
+                      ).code
+    return vary_length
+
+
+def test_normal_case_is_200(client, image):
+    response = client.POST('/v1', body=image, content_type='image/jpeg')
+    assert response.code == 200
+    assert filenames(zipfile(response.body)) == ['160', '48']
+
+def test_length_too_great_is_413(vary_length):
+    assert vary_length(262145) == 413
+
+def test_length_overstated_is_400(vary_length):
+    assert vary_length(262144) == 400
+    assert vary_length(20888) == 400
+
+def test_length_properly_stated_is_200(vary_length):
+    assert vary_length(20887, False) == 200
+
+def test_length_understated_is_400(vary_length):
+    assert vary_length(20886) == 400
+    assert vary_length(1) == 400
+    assert vary_length(0) == 400
+    assert vary_length(-1) == 400

www/v1.spt

@@ -6,20 +6,59 @@ from aspen import Response
 from PIL import Image
 
 warnings.simplefilter('error', Image.DecompressionBombWarning)
+
+
+MAX_SIZE = 256 * 1024
+
+
+class CappedStringIO(object):
+    def __init__(self, fp, target):
+        self.io = StringIO(fp)
+        self.seek(0)
+        self.__nread = 0
+        self.__target = target
+
+    def read(self, *a, **kw):
+        out = self.io.read(*a, **kw)
+        self.__nread += len(out)
+        if self.__nread > self.__target:
+            raise Response(400)
+        return out
+
+    def reset_nread(self):
+        self.__nread = 0
+
+    def check_nread(self):
+        if self.__nread != self.__target:
+            raise Response(400)
+
+    def seek(self, *a, **kw): self.io.seek(*a, **kw)
+    def tell(self, *a, **kw): self.io.tell(*a, **kw)
+
 [---]
 request.allow('POST')
 
-if int(request.headers['Content-Length']) > 256 * 1024:
+length = request.headers['Content-Length']
+try:
+    length = int(length)
+except ValueError:
+    raise Response(400)
+if length > 256 * 1024:
     raise Response(413)
 
 image_type = request.headers['Content-Type']
 if image_type not in ('image/png', 'image/jpeg'):
     raise Response(415)
 
 # Load the image.
-fp = StringIO(request.raw_body)
-fp.seek(0)
-image = Image.open(fp)
+fp = CappedStringIO(request.raw_body, length)
+try:
+    image = Image.open(fp)
+    fp.reset_nread()    # open reads *some* data
+    image.load()        # load reads *all* data ...
+except IOError:         # ... and both can raise IOError for parse errors
+    raise Response(400)
+fp.check_nread()        # still need this tho
 
 # Crop to a square.
 w, h = image.size