How can I get a non-blocking SSL Socket in Asyncio? - python-asyncio

So, I can pretty easily create an SSLSocket in asyncio:
# Create a Socket
sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
sock.setblocking(False)
# Connect
await loop.sock_connect(sock, host_address) # type: ignore
# Create an SSL Context
ssl_context = ssl.create_default_context(cafile=PROXY_CA_BUNDLE)
...
# Wrap the SSL Context Around the Socket
def do_handshake(loop, sock, waiter):
sock_fd = sock.fileno()
try:
sock.do_handshake()
except ssl.SSLWantReadError:
loop.remove_reader(sock_fd)
loop.add_reader(sock_fd, do_handshake,
loop, sock, waiter)
return
except ssl.SSLWantWriteError:
loop.remove_writer(sock_fd)
loop.add_writer(sock_fd, do_handshake,
loop, sock, waiter)
return
loop.remove_reader(sock_fd)
loop.remove_writer(sock_fd)
waiter.set_result(None)
waiter = loop.create_future()
do_handshake(loop, sslconn, waiter)
await waiter
The problem is that an SSLSocket breaks the non-blocking interface of the original socket library, so it is no longer compatible with other asyncio methods like asyncio.sock_sendall. Is there a way to instead wrap a socket and have it still respect the original socket interface?

Related

Setting the call order of mutliple asyncio.Protocol.data_received

When using asyncio for receiving multiple files over a TCP socket I struggle with the call order of received_data. When sending 3 data streams at once my output is the following:
DEBUG connection_made ('127.0.0.1', 33972) new connection was made
DEBUG connection_made ('127.0.0.1', 33974) new connection was made
DEBUG connection_made ('127.0.0.1', 33976) new connection was made
DEBUG data_received ('127.0.0.1', 33976) data received from new connection
DEBUG data_received ('127.0.0.1', 33974) data received from new connection
DEBUG data_received ('127.0.0.1', 33972) data received from new connection
I assume that its behaviour is analog to a stack where the data is received from the newest to the oldest connection made, but this is only a guess.
Is it possible to change that behaviour in a way that the data is received in the order that the connections were made? This is important because I need the data received from the first connection for further processing the following connections.
My code is the following:
import asyncio
class AIO(asyncio.Protocol):
def __init__(self):
self.extra = bytearray()
def connection_made(self, transport):
global request_time
peer = transport.get_extra_info('peername')
logger.debug('{} new connection was made'.format(peer))
self.transport = transport
request_time = str(time.time())
def data_received(self, request, msgid=1):
peer = self.transport.get_extra_info('peername')
logger.debug('{} data received from new connection'.format(peer))
self.extra.extend(request)
First, I would recommend using the higher-level streams API instead of the transport/protocols. Then if you need to maintain the order observed when connections were made, you can enforce it yourself using a series of asyncio.Events. For example:
import asyncio, itertools
class Comm:
def __init__(self):
self.extra = bytearray()
self.preceding_evt = None
async def serve(self, r, w):
preceding_evt = self.preceding_evt
self.preceding_evt = this_evt = asyncio.Event()
while True:
request = await r.read(4096)
# wait for the preceding connection to receive and store data
# before we store ours
if preceding_evt is not None:
await preceding_evt.wait()
preceding_evt.clear()
# self.extra now contains data from previous connections
self.extra.extend(request)
# inform the subsequent connection that we got data
this_evt.set()
w.close()
async def main():
comm = Comm()
server = await asyncio.start_server(comm.serve, '127.0.0.1', 8888)
async with server:
await server.serve_forever()
asyncio.run(main())

How do I avoid the loop argument

The following code is part of some automated tests that I have written in python 3.6:
connected = False
def aiohttp_server(loop):
async def handler(msg, session):
global connected
if msg.type == sockjs.MSG_OPEN:
connected = True
if msg.type == sockjs.MSG_CLOSE:
connected = False
app = web.Application(loop=loop)
sockjs.add_endpoint(app, handler)
runner = web.AppRunner(app)
return runner
def run_server(runner, loop):
logging.basicConfig(level=logging.DEBUG,
format='%(asctime)s %(levelname)s %(message)s')
asyncio.set_event_loop(loop)
loop.run_until_complete(runner.setup())
site = web.TCPSite(runner, 'localhost', 8080)
loop.run_until_complete(site.start())
loop.run_forever()
def start_server():
loop = asyncio.new_event_loop()
t = threading.Thread(target=run_server, args=(aiohttp_server(loop),loop,), daemon=True)
t.start()
time.sleep(0.01)
Basically, calling start_server should initiate a simple web server with a sockjs endpoint named /sockjs
I am not yet a master of python's async keyword. There are two issues, that I suspect are related:
Firstly, I am getting a deprecation warning on the app = web.Application(loop=loop) statement:
/home/peter/incubator/sockjs_client/tests/test_sockjs_client.py:25: DeprecationWarning: loop argument is deprecated
app = web.Application(loop=loop)
/home/peter/.local/lib/python3.6/site-packages/sockjs/route.py:54: DeprecationWarning: loop property is deprecated
manager = SessionManager(name, app, handler, app.loop)
And secondly, the tests fail occasionally. I believe that, depending on machine load, sometimes the server hasn't had enough time to start before the test code actually starts executing.
Basically, what I need is for the start_server function to initialise a web application with a websocket endpoint, and not return until the application is prepared to accept websocket connections.
Firstly, I am getting a deprecation warning on the app = web.Application(loop=loop) statement:
The recommended way to avoid passing around the loop everywhere is to switch to asyncio.run. Instead of managing the loop manually, let asyncio.run create (and close) the loop for you. If all your work is done in coroutines, you can access the loop with get_event_loop() or get_running_loop().
Basically, what I need is for the start_server function to initialise a web application with a websocket endpoint, and not return until the application is prepared to accept websocket connections.
You can pass a threading.Event to the thread that gets set when the site is set up, and wait for it in the main thread.
Here is an (untested) example that implements both suggestions:
connected = False
def aiohttp_server():
async def handler(msg, session):
global connected
if msg.type == sockjs.MSG_OPEN:
connected = True
if msg.type == sockjs.MSG_CLOSE:
connected = False
app = web.Application()
sockjs.add_endpoint(app, handler)
return web.AppRunner(app)
async def run_server(ready):
logging.basicConfig(level=logging.DEBUG,
format='%(asctime)s %(levelname)s %(message)s')
runner = aiohttp_server()
await runner.setup()
site = web.TCPSite(runner, 'localhost', 8080)
await site.start()
ready.set()
# emulates loop.run_forever()
await asyncio.get_running_loop().create_future()
def start_server():
ready = threading.Event()
threading.Thread(target=asyncio.run, args=(aiohttp_server(ready),),
daemon=True).start()
ready.wait()
Please upgrade sockjs to the newest version.
It doesn't require passing the loop anymore.

NATS does not raise an exception when disconnecting

I use an almost standard example of using NATS on python asyncio. I want to receive a message, process it, and send the result
back to the queue, but when the NATS is disconnected (e.g. reboot gnats), the exception does not raise. I even did await asyncio.sleep (1, loop = loop) to change the context
and the disconnect -> reconnect exception was thrown, but this does not happen. What am I doing wrong? May be it's a bug?
import asyncio
from nats.aio.client import Client as NATS
import time
async def run(loop):
nc = NATS()
await nc.connect(io_loop=loop)
async def message_handler(msg):
subject = msg.subject
reply = msg.reply
data = msg.data.decode()
print("Received a message on '{subject} {reply}': {data}".format(
subject=subject, reply=reply, data=data))
# Working
time.sleep(10)
# If nats disconnects at this point, the exception will not be caused
# and will be made attempt to send a message by nc.publish
await asyncio.sleep(2, loop=loop)
print("UNSLEEP")
await nc.publish("test", "test payload".encode())
print("PUBLISHED")
# Simple publisher and async subscriber via coroutine.
await nc.subscribe("foo", cb=message_handler)
while True:
await asyncio.sleep(1, loop=loop)
await nc.close()
if __name__ == '__main__':
loop = asyncio.get_event_loop()
loop.run_until_complete(run(loop))
loop.close()
NATs is built on top of TCP.
TCP has no reliable disconnection signal by definition.
To solve it any messaging system should use a kind of ping messages and drop the connection if timeout occurs.
Strictly speaking you will get disconnection event sometimes, but it may take up to 2 hours (depends on your OS settings).

How do I use asyncio.sslproto.SSLProtocol?

I am trying to talk to an ElectrumX server using JSON-RPC over TLS, but through Tor (SOCKS proxy on localhost).
When running the following code, drain() in asyncio.streams is calling _drain_helper in SSLProtocol, which I don't know how to implement.
If I just make it a no-op, it seems to not work.
I know that the JSON-RPC server is working because I have tested with
echo -ne '{"id":0,"args":["3.0.2","1.1"],"method":"server.version"}\n' | socat stdio openssl-connect:songbird.bauerj.eu:50002,verify=0
My attempt at using TLS through SOCKS in Python with asyncio:
from asyncio.sslproto import SSLProtocol
import aiosocks
import asyncio
loop = None
class MySSLProtocol(SSLProtocol):
def __init__(otherself):
super().__init__(loop, None, False , None)
# app_proto context waiter
async def l(fut):
try:
socks4addr = aiosocks.Socks4Addr("127.0.0.1", 9050)
transport, protocol = await aiosocks.create_connection(MySSLProtocol, proxy=socks4addr, proxy_auth=None, dst=("songbird.bauerj.eu", 50002))
reader = asyncio.StreamReader()
reader.set_transport(transport)
writer = asyncio.StreamWriter(transport, protocol, reader, loop)
writer.write(b'{"id":0,"method":"server.version","args":["3.0.2", "1.1"]}\n')
await writer.drain()
print(await reader.readuntil(b"\n"))
fut.set_result("finished")
except BaseException as e:
fut.set_exception(e)
def f():
global loop
loop = asyncio.get_event_loop()
fut = asyncio.Future()
asyncio.ensure_future(l(fut))
loop.run_until_complete(fut)
print(fut.result())
loop.close()
f()
I am using aiosocks from master. Commit 932374c
asyncio.sslproto and SSLProtocol are part of asyncio private API.
You should never use the class directly or derive from it.
For working with SSL please pass normal protocol (derived from asyncio.Protocol) and ssl.SSLContext as ssl param into loop.create_connection() / loop.create_server().
I'm using it like so:
class HTTP(asyncio.Protocol):
def __init__(self, config: Config):
self.config = config
def data_received(self, data) -> None:
print(data)
class HTTPS:
def __new__(cls, config: Config):
ssl_context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
ssl_context.load_cert_chain(config.rsa_cert, config.rsa_key)
return asyncio.sslproto.SSLProtocol(
loop=asyncio.get_running_loop(),
app_protocol=HTTP(config),
sslcontext=ssl_context,
waiter=None,
server_side=True,
)
See mitm for working example. Not the way I would've chosen to use it, but the only way I could figure out. From my knowledge there is no way to upgrade an asyncio.Protocol to use SSL - you must create a new one.

pyzmq recv_json can't decode message sent by send_json

Here is my code with the extraneous stuff stripped out:
coordinator.py
context = zmq.Context()
socket = context.socket(zmq.ROUTER)
port = socket.bind_to_random_port(ZMQ_ADDRESS)
poller = zmq.Poller()
poller.register(socket, zmq.POLLIN)
while True:
event = poller.poll(1)
if not event:
continue
process_id, val = socket.recv_json()
worker.py
context = zmq.Context()
socket = context.socket(zmq.DEALER)
socket.connect('%s:%s' % (ZMQ_ADDRESS, kwargs['zmq_port']))
socket.send_json(
(os.getpid(), True)
)
what happens when I run it:
process_id, val = socket.recv_json()
File "/Users/anentropic/.virtualenvs/myproj/lib/python2.7/site-packages/zmq/sugar/socket.py", line 380, in recv_json
return jsonapi.loads(msg)
File "/Users/anentropic/.virtualenvs/myproj/lib/python2.7/site-packages/zmq/utils/jsonapi.py", line 71, in loads
return jsonmod.loads(s, **kwargs)
File "/Users/anentropic/.virtualenvs/myproj/lib/python2.7/site-packages/simplejson/__init__.py", line 451, in loads
return _default_decoder.decode(s)
File "/Users/anentropic/.virtualenvs/myproj/lib/python2.7/site-packages/simplejson/decoder.py", line 406, in decode
obj, end = self.raw_decode(s)
File "/Users/anentropic/.virtualenvs/myproj/lib/python2.7/site-packages/simplejson/decoder.py", line 426, in raw_decode
raise JSONDecodeError("No JSON object could be decoded", s, idx)
JSONDecodeError: No JSON object could be decoded: line 1 column 0 (char 0)
and if I dig in with ipdb:
> /Users/anentropic/.virtualenvs/myproj/lib/python2.7/site-packages/zmq/sugar/socket.py(380)recv_json()
379 msg = self.recv(flags)
--> 380 return jsonapi.loads(msg)
381
ipdb> p msg
'\x00\x9f\xd9\x06\xa2'
hmm, that doesn't look like JSON... is this a bug in pyzmq? am I using it wrong?
Hmm, ok, found the answer.
There is an annoying asymmetry in the ØMQ interface, so you have to be aware of the type of socket you are using.
In this case my use of ROUTER/DEALER architecture means that the JSON message sent from the DEALER socket, when I do send_json, gets wrapped in multipart message envelope. The first part is a client id (I guess this is the '\x00\x9f\xd9\x06\xa2' that I got above) and the second part is the JSON string we are interested in.
So in the last line of my coordinator.py I need to do this instead:
id_, msg = socket.recv_multipart()
process_id, val = json.loads(msg)
IMHO this is bad design on the part of ØMQ/pyzmq, the library should abstract this away and have just send and recv methods, that just work.
I got the clue from this question How can I use send_json with pyzmq PUB SUB so it looks like PUB/SUB architecture has the same issue, and no doubt others too.
This is described in the docs but it's not very clear
http://zguide.zeromq.org/page:all#The-Asynchronous-Client-Server-Pattern
Update
In fact, I found in my case I could simplify the code further, by making use of the 'client id' part of the message envelope directly. So the worker just does:
context = zmq.Context()
socket = context.socket(zmq.DEALER)
socket.identity = str(os.getpid()) # or I could omit this and use ØMQ client id
socket.connect('%s:%s' % (ZMQ_ADDRESS, kwargs['zmq_port']))
socket.send_json(True)
It's also worth noting that when you want to send a message the other direction, from the ROUTER, you have to send it as multipart, specifying which client it is destined for, eg:
coordinator.py
context = zmq.Context()
socket = context.socket(zmq.ROUTER)
port = socket.bind_to_random_port(ZMQ_ADDRESS)
poller = zmq.Poller()
poller.register(socket, zmq.POLLIN)
pids = set()
while True:
event = poller.poll(1)
if not event:
continue
process_id, val = socket.recv_json()
pids.add(process_id)
# need some code in here to decide when to stop listening
# and break the loop
for pid in pids:
socket.send_multipart([pid, 'a string message'])
# ^ do your own json encoding if required
I guess there is probably some ØMQ way of doing a broadcast message rather than sending to each client in a loop as I do above. I wish the docs just had a clear description of each available socket type and how to use them.

Resources