Merge pull request #756 from tumb1er/fix_multipart_chunk

Fix multipart chunk

aiohttp/multipart.py

@@ -305,20 +305,24 @@ def _read_chunk_from_stream(self, size):
         """
         assert size >= len(self._boundary) + 2, \
             'Chunk size must be greater or equal than boundary length + 2'
-        if self._prev_chunk is None:
+        first_chunk = self._prev_chunk is None
+        if first_chunk:
             self._prev_chunk = yield from self._content.read(size)
 
         chunk = yield from self._content.read(size)
 
         window = self._prev_chunk + chunk
         sub = b'\r\n' + self._boundary
-        idx = window.find(sub, len(self._prev_chunk) - len(sub))
+        if first_chunk:
+            idx = window.find(sub)
+        else:
+            idx = window.find(sub, len(self._prev_chunk) - len(sub))
         if idx >= 0:
             # pushing boundary back to content
             self._content.unread_data(window[idx:])
             if size > idx:
                 self._prev_chunk = self._prev_chunk[:idx]
-            chunk = window[size:idx]
+            chunk = window[len(self._prev_chunk):idx]
             if not chunk:
                 self._at_eof = True
         result = self._prev_chunk

aiohttp/streams.py

@@ -159,8 +159,6 @@ def wait_eof(self):
     def unread_data(self, data):
         """ rollback reading some data from stream, inserting it to buffer head.
         """
-        assert not self._eof, 'unread_data after feed_eof'
-
         if not data:
             return
 

tests/test_multipart.py

@@ -168,6 +168,29 @@ def test_read_chunk_without_content_length(self):
         self.assertEqual(c1 + c2, b'Hello, world!')
         self.assertEqual(c3, b'')
 
+    def test_read_incomplete_chunk(self):
+        stream = Stream(b'')
+
+        def prepare(data):
+            f = asyncio.Future(loop=self.loop)
+            f.set_result(data)
+            return f
+
+        with mock.patch.object(stream, 'read', side_effect=[
+            prepare(b'Hello, '),
+            prepare(b'World'),
+            prepare(b'!\r\n--:'),
+            prepare(b'')
+        ]):
+            obj = aiohttp.multipart.BodyPartReader(
+                self.boundary, {}, stream)
+            c1 = yield from obj.read_chunk(8)
+            self.assertEqual(c1, b'Hello, ')
+            c2 = yield from obj.read_chunk(8)
+            self.assertEqual(c2, b'World')
+            c3 = yield from obj.read_chunk(8)
+            self.assertEqual(c3, b'!')
+
     def test_multi_read_chunk(self):
         stream = Stream(b'Hello,\r\n--:\r\n\r\nworld!\r\n--:--')
         obj = aiohttp.multipart.BodyPartReader(self.boundary, {}, stream)

tests/test_streams.py

@@ -409,6 +409,11 @@ def test_unread_data(self):
         data = self.loop.run_until_complete(stream.read(4))
         self.assertEqual(b'line', data)
 
+        stream.feed_eof()
+        stream.unread_data(b'at_eof')
+        data = self.loop.run_until_complete(stream.read(6))
+        self.assertEqual(b'at_eof', data)
+
     def test_exception(self):
         stream = self._make_one()
         self.assertIsNone(stream.exception())