aio-libs / aiohttp

Asynchronous HTTP client/server framework for asyncio and Python
https://docs.aiohttp.org
Other
15.17k stars 2.02k forks source link

WebSocket prepare property incorrectly reflects prepared state #6009

Open balloob opened 3 years ago

balloob commented 3 years ago

Describe the bug

To Reproduce

wsock = web.WebSocketResponse()
try:
    # Bug can happen when this job times out during `WebSocketResponse `calling `super().prepare()`
    # and `StreamResponse` is inside `await self._start` and client connection is dropped but OS didn't notice
    async with async_timeout(10):
        await wsock.prepare()
except asyncio.TimeoutError:
    if wsock.prepared:
        # This raises because `close()` checks if `self._writer` is not None,
        # which is set inside `StreamResponse._post_start`
        await wsock.close()

Expected behavior

If wsock.prepared returns True, it should be safe to close the connection.

Logs/tracebacks

File "server/api.py", line 264, in handle
    await wsock.close()
  File "aiohttp/web_ws.py", line 329, in close
    raise RuntimeError("Call .prepare() first")

### Python Version

```console
$ python --version
3.9

aiohttp Version

$ python -m pip show aiohttp
3.7.4.post0

multidict Version

$ python -m pip show multidict
5.1.0

yarl Version

$ python -m pip show yarl
1.6.3

OS

Debian

Related component

Server

Additional context

No response

Code of Conduct

bdraco commented 11 months ago

Full reproducer

from aiohttp import web
from async_timeout import timeout
import asyncio
from multidict import CIMultiDict, CIMultiDictProxy, istr
from aiohttp.test_utils import make_mocked_request
from unittest import mock
import aiosignal
from functools import partial

def make_request(method, path, headers=None, protocols=False):
    app = ret = mock.Mock()
    app.loop = loop
    app._debug = False
    app.on_response_prepare = aiosignal.Signal(ret)
    app.on_response_prepare.freeze()

    protocol = mock.Mock()
    protocol.set_parser.return_value = ret

    if headers is None:
        headers = CIMultiDict(
            {
                "HOST": "server.example.com",
                "UPGRADE": "websocket",
                "CONNECTION": "Upgrade",
                "SEC-WEBSOCKET-KEY": "dGhlIHNhbXBsZSBub25jZQ==",
                "ORIGIN": "http://example.com",
                "SEC-WEBSOCKET-VERSION": "13",
            }
        )
    if protocols:
        headers["SEC-WEBSOCKET-PROTOCOL"] = "chat, superchat"

    req = make_mocked_request(
        method, path, headers, app=app, protocol=protocol, loop=app.loop
    )

    return req

async def repro():
    request = make_request("GET", "/ws")
    wsock = web.WebSocketResponse()

    # Simulate `WebSocketResponse._start` taking a long time
    async def _prepare_headers(*args, **kwargs):
        await asyncio.sleep(10)

    wsock._prepare_headers = _prepare_headers
    try:
        # Bug can happen when this job times out during `WebSocketResponse `calling `super().prepare()`
        # and `StreamResponse` is inside `await self._start` and client connection is dropped but OS didn't notice
        async with timeout(3):
            await wsock.prepare(request)
    except asyncio.TimeoutError:
        assert wsock.prepared
        if wsock.prepared:
            # This raises because `close()` checks if `self._writer` is not None,
            # which is set inside `StreamResponse._post_start`
            await wsock.close()

async def run():
    await repro()

if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    loop.run_until_complete(run())