Context information storage for asyncio
Tags: dev, python, opensource
Published: 2018-03-01 23:22:27

---
title: "Context information storage for asyncio"
description: Implement context information storage for asyncio in Python
published: true
tags: Dev,Python, opensource 
canonical_url: https://blog.sqreen.io/asyncio/
cover_image: https://thepracticaldev.s3.amazonaws.com/i/fxhwdc0bis95nylm2jwk.jpg
---

At [Sqreen](https://www.sqreen.io/), we are building an agent based on dynamic instrumentation. It detects security incidents from inside an application (injections, cross-site scripting etc.) and allows users to configure actions (blocking the attack, logging a stack trace etc.) without requiring code modification. The mechanisms behind dynamic instrumentation in Python are described in [a previous blog post](https://blog.sqreen.io/dynamic-instrumentation-agent-for-python/). Dynamic instrumentation is also used in Application Performance Management (APM) solutions, such as [Datadog](https://www.datadoghq.com/), [Instana](https://www.instana.com/) and [New Relic](https://newrelic.com/).

Instrumenting the code allows us to execute callbacks before calling potentially hazardous functions. For example, to [protect against SQL injections](https://blog.sqreen.io/block-sql-injections-not-customers/), we transparently wrap the method Cursor.execute with a security layer:

```python
def sqreen_execute(self, sql_stmt, *sql_params):
    # Before executing the SQL statement, check it was not built with
    # malicious, unescaped request parameters in it.
    if has_malicious_param(sql_stmt, request.params):
        # If there is, this is an SQL injection. Abort!
        raise SQLinjection(remote_addr=request.remote_addr)
    else:
        # If not, we can safely call the original method.
        return self.execute(sql_stmt, *sql_params)
```

Let’s assume the instrumented code contains a vulnerable pattern like:

```python
@app.route('/posts')
def posts(request):
    sql_stmt = 'SELECT * FROM posts WHERE id=%s' % request.params['id']

    # With dynamic instrumentation, sqreen_execute is transparently called
    # instead of cursor.execute.
    posts = cursor.execute(sql_stmt)

    return posts_template.render(posts)
```

Then the nominal request /posts/?id=42 will be executed (although unescaped, the request parameter is not malicious) but the malicious request /posts/?id=42 OR 1=1 won’t be. This way, we are able to protect the app without breaking it!

## Context information storage

As we’ve seen above, the function sqreen\_execute needs to know the current request to test the SQL statement safety. How can it get it?

Since our function transparently replaces Cursor.execute, it needs to have the same signature, hence we can’t pass the request as a parameter to sqreen\_execute. Some web frameworks provide functions to get the current requests, but not all of them, and we strive for a universal solution.

What we can do is to insert a middleware (or instrument the framework’s request handling mechanism) to store the current request in a global variable:

```python
CURRENT_REQUEST = None

def set_request(request):
    global CURRENT_REQUEST
    current_request = request

def get_request():
    return CURRENT_REQUEST
```

But there is a catch: web frameworks are able to handle several requests concurrently, for obvious performance reasons. So the above pattern won’t work: we may receive a first request request\_1 (and store it in CURRENT\_REQUEST), and before serving it receive a second request request\_2 (overriding request\_1 in CURRENT\_REQUEST). At the time we look for SQL injections in request\_1, we will mess it up with request\_2! So we need a stronger, concurrent-safe mechanism to store the current request.

### Thread-local storage

To tackle this issue, we need to know how concurrency is implemented. Most of Python web frameworks use threads: this is notably the case of [Django](https://www.djangoproject.com/), [Flask](http://flask.pocoo.org/) and [Pyramid](https://trypyramid.com/), which are probably the most popular. They implement a common communication protocol with web servers, called WSGI (Web Server Gateway Interface) and initially described in [PEP 333](https://www.python.org/dev/peps/pep-0333/).

WSGI servers also use threads, along with processes, to spawn several application instances. Multiprocessing is not an issue here, since each process will handle its own copy of CURRENT\_REQUEST. So, we just have to find a solution to let a service thread to store the request it is currently dealing with without impacting other threads.

And Python offers a solution for that. The function threading.local [in the standard library](https://docs.python.org/3/library/threading.html#threading.local) return a namespace object whose values are thread specific. This allows us to implement thread-safe request storage as follows:

```python
import threading

RUNTIME_STORAGE = threading.local()
RUNTIME_STORAGE.request = None

def set_request(request):
    RUNTIME_STORAGE.request = request

def get_request():
    return RUNTIME_STORAGE.request
```

### What about asyncio?

In Python 3.4, a new concurrency model was introduced: [asyncio](https://docs.python.org/3/library/asyncio.html). It provides infrastructure for single-threaded, asynchronous programming, including:

- Coroutine functions, defined with async def, whose execution can be paused using the await keyword with another coroutine, and resumed once the other coroutine is completed.
- An event loop to schedule and execute coroutines.

Here is an example of asynchronous code.

```python
import asyncio

async def compute(x, y):
    print("Compute %s + %s ..." % (x, y))
    await asyncio.sleep(1.0)
    return x + y

async def print_sum(x, y):
    result = await compute(x, y)
    print("%s + %s = %s" % (x, y, result))

loop = asyncio.get_event_loop()
loop.run_until_complete(print_sum(1, 2))
loop.close()
```

There are two coroutine functions, print\_sum and compute. At execution time

- The event loop enters print\_sum and immediately hands over to compute.
- compute prints the computation and hands over to asyncio.sleep.
- Nothing is done in the next second. If other tasks were scheduled in the event loop, they could be executed in the meantime, something that is not possible with the blocking function time.sleep.
- compute is resumed and completed.
- print\_sum is resumed and completed.

asyncio is a great model for concurrency when IO is involved: when the code being executed is blocked waiting for an answer (for example, DB results), the program can switch to other tasks and come back to it later. It is less system-expensive than threads, and is usually faster when slow IO operations are involved.

This makes asyncio well suited for network operations and, despite being relatively young, [several web frameworks](https://github.com/timofurrer/awesome-asyncio#web-frameworks) have been developed around it. Among them, we recently brought support for [aiohttp](https://aiohttp.readthedocs.io/) in our agent. This was a very interesting and challenging task since we had no support for aiohttp at all so far, and an important issue we met was with the request storage mechanism.

Here is what can happen: we receive a first request request\_1 and start dealing with it in a coroutine. At some point, the coroutine is suspended and the event loop hands over to another one that handles request\_2. The important point is that these two coroutines are executed in the _same thread_, so threading.local contains the same data for both. When the first coroutine resumes, RUNTIME\_STORAGE.request has been set to request\_2: that is precisely what we want to prevent.

## First attempt: let’s use tasks!

What we need is a mechanism similar to threading.local that works with asyncio, i.e. lets us store context variables and keep track of values per asynchronous execution context.

Unfortunately, there is currently no built-in mechanism in Python to handle this. Different proposals have been made to provide a generic solution in future versions of Python ([PEP 550](https://www.python.org/dev/peps/pep-0550/), [PEP 567](https://www.python.org/dev/peps/pep-0567/)), but in the meantime, we have to devise a solution on our own.

Let’s dig a bit further into the internals of asyncio. A coroutine whose execution is scheduled is wrapped into an asyncio.Task [object](https://docs.python.org/3/library/asyncio-task.html#task), responsible for executing the coroutine object in an event loop.

![https://docs.python.org/3/_images/tulip_coro.png](https://i0.wp.com/blog.sqreen.io/wp-content/uploads/2018/02/https-docs-python-org-3-_images-tulip_coro-png.png?w=840&ssl=1)

_Sequence diagram of the example_

There is also a function asyncio.Task.current\_task that returns the currently running task. Mmh… We could use this to map the current request to the task handling it. Something like this could work:

```python
import asyncio

TASK_REQUESTS = {}

def set_request(request):
    task = asyncio.Task.current_task()
    TASK_REQUESTS[id(task)] = request

def get_request():
    task = asyncio.Task.current_task()
    return TASK_REQUESTS.get(id(task))
```

With this implementation, we’d also need a mechanism to ensure request deletion once the task is completed, to avoid accumulating old requests and cause a memory leak. A way to avoid dealing with it is to store the request _within_ the Task object, as an extra attribute:

```python
def set_request(request):
     task = asyncio.Task.current_task()
     setattr(task, 'current_request', request)
 
 def get_request():
     task = asyncio.Task.current_task()
     return getattr(task, 'current_request', None)
```

So, does it work? Let’s test!

```python
import random

class Request:
    # Dummy request object, for the sake of testing.
    pass

async def handle_request(request):
    set_request(request)
    await asyncio.sleep(random.uniform(0, 2))
    await check_request(request)

async def check_request(request):
    # Check that the stored request corresponds to the current request. If not,
    # an AssertionError is raised and the test is interrupted with an error.
    assert get_request() is request, "wrong request"

NUM_REQUESTS = 1000

loop = asyncio.get_event_loop()
coros = [handle_request(Request()) for _ in range(NUM_REQUESTS)]
loop.run_until_complete(asyncio.gather(*coros))
loop.close()
print("Success!")
```

This test simulates one thousand concurrent requests. Each one is handled in a dedicated coroutine handle\_request. This function stores the request, then pauses for a random duration (this simulates an async operation such as a DB access, and ensures that the coroutine execution flow is interleaved). When resumed, a nested coroutine check\_request is called that ensures that get\_request returns the correct request. If not, the test is interrupted by an error.

And here are some good news: the test runs smoothly with task-based request storage. It also fails with thread-local storage, which was expected but shows the test is relevant. So, have we solved our problem?

## Context inheritance between tasks

Let’s try something a bit more twisted:

```python
async def handle_request(request):
    set_request(request)
    await asyncio.gather(
        asyncio.sleep(random.uniform(0, 2)),
        check_request(request),
    )
```

Instead of executing asyncio.sleep and check\_request sequentially, this version of handle\_request runs them concurrently. This should not be a big deal: the code is a bit more concurrent, but it does not impact request handling. In particular, check\_request is still called _after_ set\_request for each request.

Nevertheless, this new test fails! Something went wrong when we introduced asyncio.gather, but what?

Well, remember that scheduled coroutines are wrapped into tasks? That’s exactly what happens here: asyncio.gather creates tasks around the arguments asyncio.sleep() and check\_request() and these tasks are executed by the event loop.

```python
async def handle_request(request):
    set_request(request)                              # Running in task 1.
    await asyncio.gather(
        asyncio.sleep(random.uniform(0, 2)),          # Create child task 2.
        check_request(request),                       # Create child task 3.
    )

async def check_request(request):
    assert get_request() is request, "wrong request"  # Running in task 3.
```

The consequence is that set\_request and get\_request are called in different tasks, making the test fail. This is not due to request interleaving, as we can check by setting NUM\_REQUESTS to 1: the test keeps failing.

In fact, when calling get\_request from a child task, we need a mechanism to retrieve the request from the parent task if it is not defined in the child task. But asyncio does not allow us to access the parent task, so this is not going to work.

On the other hand, something asyncio let us do is to replace the function called to create new tasks, a.k.a. the _task factory_. This function is called in the context of the parent task, and returns a fresh child task. Well, let’s use it to decorate the child task with the current request!

Here is what a “request-aware” task factory would look like:

```python
def request_task_factory(loop, coro):
    # This is the default way to create a child task.
    child_task = asyncio.tasks.Task(coro, loop=loop)

    # Retrieve the request from the parent task...
    parent_task = asyncio.Task.current_task(loop=loop)
    current_request = getattr(parent_task, 'current_request', None)

    # ...and store it in the child task too.
    setattr(child_task, 'current_request', current_request)

    return child_task
```

To install the task factory, we also need to call loop.set\_task\_factory(request\_task\_factory) before running the loop. So, here is the final version of our code:

```python
import asyncio
import random

class Request:
    pass

def set_request(request):
    task = asyncio.Task.current_task()
    setattr(task, 'current_request', request)

def get_request():
    task = asyncio.Task.current_task()
    return getattr(task, 'current_request', None)

def request_task_factory(loop, coro):
    child_task = asyncio.tasks.Task(coro, loop=loop)
    parent_task = asyncio.Task.current_task(loop=loop)
    current_request = getattr(parent_task, 'current_request', None)
    setattr(child_task, 'current_request', current_request)
    return child_task

async def handle_request(request):
    set_request(request)
    await asyncio.gather(
        asyncio.sleep(random.uniform(0, 2)),
        check_request(request),
    )

async def check_request(request):
    assert get_request() is request

NUM_REQUESTS = 1000

loop = asyncio.get_event_loop()
loop.set_task_factory(request_task_factory)
coros = [handle_request(Request()) for _ in range(NUM_REQUESTS)]
loop.run_until_complete(asyncio.gather(*coros))
loop.close()
```

And it works flawlessly!

## And now what?

We now have the foundations to solve the request storage issue in our agent. Since we want the agent to work as transparently as possible and not require code modification from the user, there are still two minor issues to be tackled:

- We want to automatically set up the task factory. This will be done with dynamic instrumentation.
- But if a custom task is set up by the user, we don’t want to overwrite it. Instead, we will wrap it up on our own.

Let’s start with the second problem. We can define a generic function wrap\_request\_task\_factory that takes a task factory as argument and returns a variant of it that supports request propagation. The code of the wrapped function is really close to the one of request\_task\_factory above:

```python
from functools import wraps

def wrap_request_task_factory(task_factory):

    @wraps(task_factory)
    def wrapped(loop, coro):
        child_task = task_factory(loop, coro)
        parent_task = asyncio.Task.current_task(loop=loop)
        current_request = getattr(parent_task, 'current_request', None)
        setattr(child_task, 'current_request', current_request)
        return child_task

    return wrapped
```

Then, the definition of request\_task\_factory can be simplified to:

```python
@wrap_request_task_factory
def request_task_factory(loop, coro):
    asyncio.Task.current_task(loop=loop)
```

Time to go back to dynamic instrumentation. By hooking the import system, we can transparently replace an imported class by a custom one. So let’s define a function patch\_loop\_cls that creates a custom loop class with the desired behavior:

```python
def wrap_loop_cls(loop_cls):

    class SqreenLoop(loop_cls):

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            # We want to use request_task_factory to be the default task
            # factory.
            super().set_task_factory(request_task_factory)

        def set_task_factory(self, task_factory):
            # If the user sets up a custom task factory, let's wrap it with
            # request propagation.
            wrapped_task_factory = wrap_request_task_factory(task_factory)
            super().set_task_factory(wrapped_task_factory)

    return SqreenLoop
```

This loop class transparently replaces the base one. It uses the correct task factory by default and allows the user to change it while keeping the request management layer.

## Closing words

We have published most of this work (without the instrumentation part) in a Python library called [AioContext](https://github.com/sqreen/AioContext). It comes with generic Context objects that behave like dictionaries. It also allows to restore the original task factory if contexts are no longer needed, and stores contexts as an extra attribute of the task factory to avoid messing with the asyncio.Task class itself. The documentation is available [here](https://aiocontext.readthedocs.io/).

```python
import asyncio
import aiocontext
import random

class Request:
    pass

CONTEXT = aiocontext.Context()

async def handle_request(request):
    CONTEXT['current_request'] = request
    await asyncio.gather(
        asyncio.sleep(random.uniform(0, 2)),
        check_request(request),
    )

async def check_request(request):
    assert CONTEXT['current_request'] is request

NUM_REQUESTS = 1000

loop = asyncio.get_event_loop()
aiocontext.wrap_task_factory(loop)
CONTEXT.attach(loop)
coros = [handle_request(Request()) for _ in range(NUM_REQUESTS)]
loop.run_until_complete(asyncio.gather(*coros))
loop.close()
```

This work was strongly inspired by Manual Miranda’s blog post [From Flask to aiohttp](https://medium.com/@SkyscannerEng/from-flask-to-aiohttp-22f1ddc5dd5e) and the library [aiotask-context](https://github.com/Skyscanner/aiotask-context). We want to thank him for the great contribution.

The post [Context information storage for asyncio](https://blog.sqreen.io/asyncio/) appeared first on [Sqreen Blog | Modern Application Security](https://blog.sqreen.io) and was written by Vivien Maisonneuve.