Repository: falconry/falcon
Pull request #1692 (2020-02-29 05:03:56)

fix(asgi.BoundedStream): Mixing iteration and read() can lead to subtle errors
# Summary of Changes

This patch does three things to make BoundedStream more robust and debuggable:

1. Disallow re-iteration over the stream when an iteration is already in progress.
2. Yield anything in the buffer first, before awaiting additional body events,
   when iterating over a stream that was partially read before the iteration
   began.
3. Introduce a custom subclass of ValueError, OperationNotAllowed, that
   can be raised to give the developer a more obvious hint as to what
   went wrong.

# Related Issues

#1358 


codecov[bot]: # [Codecov](https://codecov.io/gh/falconry/falcon/pull/1692?src=pr&el=h1) Report
> Merging [#1692](https://codecov.io/gh/falconry/falcon/pull/1692?src=pr&el=desc) into [master](https://codecov.io/gh/falconry/falcon/commit/e3bf192bd9a4a89e3892c9ab7c95c51f3983bdaf?src=pr&el=desc) will **not change** coverage.
> The diff coverage is `100%`.

[![Impacted file tree graph](https://codecov.io/gh/falconry/falcon/pull/1692/graphs/tree.svg?width=650&token=zOtkfNQ3h6&height=150&src=pr)](https://codecov.io/gh/falconry/falcon/pull/1692?src=pr&el=tree)

```diff
@@          Coverage Diff           @@
##           master   #1692   +/-   ##
======================================
  Coverage     100%    100%           
======================================
  Files          47      47           
  Lines        3526    3537   +11     
  Branches      548     550    +2     
======================================
+ Hits         3526    3537   +11
```


| [Impacted Files](https://codecov.io/gh/falconry/falcon/pull/1692?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [falcon/asgi/stream.py](https://codecov.io/gh/falconry/falcon/pull/1692/diff?src=pr&el=tree#diff-ZmFsY29uL2FzZ2kvc3RyZWFtLnB5) | `100% <100%> (ø)` | :arrow_up: |
| [falcon/errors.py](https://codecov.io/gh/falconry/falcon/pull/1692/diff?src=pr&el=tree#diff-ZmFsY29uL2Vycm9ycy5weQ==) | `100% <100%> (ø)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/falconry/falcon/pull/1692?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/falconry/falcon/pull/1692?src=pr&el=footer). Last update [e3bf192...27f98a1](https://codecov.io/gh/falconry/falcon/pull/1692?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/falcon/asgi/stream.py b/falcon/asgi/stream.py
index 7dae30d..ec0022c 100644
--- a/falcon/asgi/stream.py
+++ b/falcon/asgi/stream.py
@@ -15,6 +15,9 @@
 """ASGI BoundedStream class."""
 
 
+from falcon.errors import OperationNotAllowed
+
+
 __all__ = ['BoundedStream']
 
 
@@ -87,12 +90,14 @@ class BoundedStream:
         '_buffer',
         '_bytes_remaining',
         '_closed',
+        '_iteration_started',
         '_pos',
         '_receive',
     ]
 
     def __init__(self, receive, content_length=None):
         self._closed = False
+        self._iteration_started = False
 
         self._receive = receive
         self._buffer = b''
@@ -108,9 +113,9 @@ class BoundedStream:
         self._pos = 0
 
     def __aiter__(self):
-        # NOTE(kgriffs): Technically we should be returning an async iterator
-        #   here instead of an async generator, but in practice the caller
-        #   should be happy as long as the returned object is iterable.
+        # NOTE(kgriffs): This returns an async generator, but that's OK because
+        #   it also implements the iterator protocol defined in PEP 492, albeit
+        #   in a more efficient way than a regular async iterator.
         return self._iter_content()
 
     # -------------------------------------------------------------------------
@@ -224,7 +229,7 @@ class BoundedStream:
         """
 
         if self._closed:
-            raise ValueError(
+            raise OperationNotAllowed(
                 'This stream is closed; no further operations on it are permitted.'
             )
 
@@ -298,7 +303,7 @@ class BoundedStream:
         """
 
         if self._closed:
-            raise ValueError(
+            raise OperationNotAllowed(
                 'This stream is closed; no further operations on it are permitted.'
             )
 
@@ -360,13 +365,9 @@ class BoundedStream:
 
         return data
 
-    # NOTE: In docs, tell people to not mix reading different modes - make
-    #   sure you exhaust in the finally if you are reading something
-    #   in middleware, or a chance something else might read it. Don't want someone
-    #   to end up trying to read a half-read thing anyway!
     async def _iter_content(self):
         if self._closed:
-            raise ValueError(
+            raise OperationNotAllowed(
                 'This stream is closed; no further operations on it are permitted.'
             )
 
@@ -374,9 +375,17 @@ class BoundedStream:
             yield b''
             return
 
-        # TODO(kgriffs): Should we check for any buffered data and return
-        #   that first? Or simply raise an error if any data has already
-        #   been read?
+        if self._iteration_started:
+            raise OperationNotAllowed('This stream is already being iterated over.')
+
+        self._iteration_started = True
+
+        if self._buffer:
+            next_chunk = self._buffer
+            self._buffer = b''
+
+            self._pos += len(next_chunk)
+            yield next_chunk
 
         while self._bytes_remaining > 0:
             event = await self._receive()
diff --git a/falcon/errors.py b/falcon/errors.py
index 0f13a46..630ecff 100644
--- a/falcon/errors.py
+++ b/falcon/errors.py
@@ -58,6 +58,12 @@ class UnsupportedError(RuntimeError):
     """The method or operation is not supported."""
 
 
+# NOTE(kgriffs): This inherits from ValueError to be consistent with the type
+#   raised by Python's built-in file-like objects.
+class OperationNotAllowed(ValueError):
+    """The requested operation is not allowed."""
+
+
 class HTTPBadRequest(HTTPError):
     """400 Bad Request.
 


Test patch:
diff --git a/tests/asgi/test_boundedstream_asgi.py b/tests/asgi/test_boundedstream_asgi.py
index dbc83c1..618bdf6 100644
--- a/tests/asgi/test_boundedstream_asgi.py
+++ b/tests/asgi/test_boundedstream_asgi.py
@@ -195,7 +195,15 @@ def test_read_chunks(body, chunk_size):
 
         assert b''.join(chunks) == body
 
-    for t in (test_nonmixed, test_mixed_a, test_mixed_b):
+    async def test_mixed_iter():
+        s = stream()
+
+        chunks = [await s.read(chunk_size)]
+        chunks += [data async for data in s]
+
+        assert b''.join(chunks) == body
+
+    for t in (test_nonmixed, test_mixed_a, test_mixed_b, test_mixed_iter):
         testing.invoke_coroutine_sync(t)
         testing.invoke_coroutine_sync(t)
 
@@ -219,6 +227,30 @@ def test_exhaust_with_disconnect():
     testing.invoke_coroutine_sync(t)
 
 
+def test_iteration_already_started():
+    body = testing.rand_string(1, 2048).encode()
+    s = _stream(body)
+
+    async def t():
+        stream_iter = s.__aiter__()
+
+        chunks = [await stream_iter.__anext__()]
+
+        with pytest.raises(ValueError):
+            stream_iter2 = s.__aiter__()
+            await stream_iter2.__anext__()
+
+        while True:
+            try:
+                chunks.append(await stream_iter.__anext__())
+            except StopAsyncIteration:
+                break
+
+        assert b''.join(chunks) == body
+
+    testing.invoke_coroutine_sync(t)
+
+
 def _stream(body, content_length=None):
     emitter = testing.ASGIRequestEventEmitter(body)
     return asgi.BoundedStream(emitter, content_length=content_length)