Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

How to detect write failure in asyncio?

As a simple example, consider the network equivalent of /dev/zero, below. (Or more realistically, just a web server sending a large file.)

If a client disconnects early, you get a barrage of log messages:

WARNING:asyncio:socket.send() raised exception.

But I'm not finding any way to catch said exception. The hypothetical server continues reading gigabytes from disk and sending them to a dead socket, with no effort on the client's part, and you've got yourself a DoS attack.

The only thing I've found from the docs is to yield from a read, with an empty string indicating closure. But that's no good here because a normal client isn't going to send anything, blocking the write loop.

What's the right way to detect failed writes, or be notified that the TCP connection has been closed, with the streams API or otherwise?

Code:

from asyncio import *
import logging

@coroutine
def client_handler(reader, writer):
    while True:
        writer.write(bytes(1))
        yield from writer.drain()

logging.basicConfig(level=logging.INFO)
loop = get_event_loop()
coro = start_server(client_handler, '', 12345)
server = loop.run_until_complete(coro)
loop.run_forever()
like image 652
jwelsh Avatar asked Aug 24 '15 05:08

jwelsh


2 Answers

The stream based API doesn't have a callback you can specify for when the connection is closed. But the Protocol API does, so use it instead: https://docs.python.org/3/library/asyncio-protocol.html#connection-callbacks

like image 166
Jashandeep Sohi Avatar answered Oct 17 '22 14:10

Jashandeep Sohi


This is a little bit strange, but you can actually allow an exception to reach the client_handler coroutine by forcing it to yield control to the event loop for one iteration:

import asyncio
import logging

@asyncio.coroutine
def client_handler(reader, writer):
    while True:
        writer.write(bytes(1))
        yield  # Yield to the event loop
        yield from writer.drain()

logging.basicConfig(level=logging.INFO)
loop = asyncio.get_event_loop()
coro = asyncio.start_server(client_handler, '', 12345)
server = loop.run_until_complete(coro)
loop.run_forever()

If I do that, I get this output when I kill the client connection:

ERROR:asyncio:Task exception was never retrieved
future: <Task finished coro=<client_handler() done, defined at aio.py:4> exception=ConnectionResetError(104, 'Connection reset by peer')>
Traceback (most recent call last):
  File "/usr/lib/python3.4/asyncio/tasks.py", line 238, in _step
    result = next(coro)
  File "aio.py", line 9, in client_handler
    yield from writer.drain()
  File "/usr/lib/python3.4/asyncio/streams.py", line 301, in drain
    raise exc
  File "/usr/lib/python3.4/asyncio/selector_events.py", line 700, in write
    n = self._sock.send(data)
ConnectionResetError: [Errno 104] Connection reset by peer

I'm really not quite sure why you need to explicitly let the event loop get control for the exception to get through - don't have time at the moment to dig into it. I assume some bit needs to get flipped to indicate the connection dropped, and calling yield from writer.drain() (which can short-circuit going through the event loop) in a loop is preventing that from happening, but I'm really not sure. If I get a chance to investigate, I'll update the answer with that info.

like image 31
dano Avatar answered Oct 17 '22 14:10

dano