test_cwe_404.py

import asyncio
import contextlib

import pytest

from redis.asyncio import Redis
from redis.asyncio.cluster import RedisCluster
from redis.asyncio.connection import async_timeout


@pytest.fixture
def redis_addr(request):
    redis_url = request.config.getoption("--redis-url")
    assert redis_url.startswith("redis://")
    host, port = redis_url[8:].split(":")
    return host, int(port)


async def pipe(
    reader: asyncio.StreamReader,
    writer: asyncio.StreamWriter,
    proxy: "DelayProxy",
    name="",
    event: asyncio.Event = None,
):
    while True:
        data = await reader.read(1000)
        if not data:
            break
        if event:
            event.set()
        await asyncio.sleep(proxy.delay)
        writer.write(data)
        await writer.drain()


class DelayProxy:
    def __init__(self, addr, redis_addr, delay: float):
        self.addr = addr
        self.redis_addr = redis_addr
        self.delay = delay
        self.send_event = asyncio.Event()
        self.redis_streams = None

    async def start(self):
        # test that we can connect to redis
        with async_timeout(2):
            redis_reader, redis_writer = await asyncio.open_connection(*self.redis_addr)
        redis_writer.close()
        self.server = await asyncio.start_server(self.handle, *self.addr)
        self.ROUTINE = asyncio.create_task(self.server.serve_forever())

    @contextlib.contextmanager
    def override(self, delay: float = 0.0):
        """
        Allow to override the delay for parts of tests which aren't time dependent,
        to speed up execution.
        """
        old = self.delay
        self.delay = delay
        try:
            yield
        finally:
            self.delay = old

    async def handle(self, reader, writer):
        # establish connection to redis
        redis_reader, redis_writer = await asyncio.open_connection(*self.redis_addr)
        pipe1 = asyncio.create_task(
            pipe(reader, redis_writer, self, "to redis:", self.send_event)
        )
        pipe2 = asyncio.create_task(pipe(redis_reader, writer, self, "from redis:"))
        await asyncio.gather(pipe1, pipe2)

    async def stop(self):
        # clean up enough so that we can reuse the looper
        self.ROUTINE.cancel()
        loop = self.server.get_loop()
        await loop.shutdown_asyncgens()


@pytest.mark.onlynoncluster
@pytest.mark.parametrize("delay", argvalues=[0.05, 0.5, 1, 2])
async def test_standalone(delay, redis_addr):

    # create a tcp socket proxy that relays data to Redis and back,
    # inserting 0.1 seconds of delay
    dp = DelayProxy(addr=("127.0.0.1", 5380), redis_addr=redis_addr, delay=delay * 2)
    await dp.start()

    for b in [True, False]:
        # note that we connect to proxy, rather than to Redis directly
        async with Redis(host="127.0.0.1", port=5380, single_connection_client=b) as r:

            with dp.override():
                await r.set("foo", "foo")
                await r.set("bar", "bar")

            dp.send_event.clear()
            t = asyncio.create_task(r.get("foo"))
            # Wait until the task has sent, and then some, to make sure it has
            # settled on the read.
            await dp.send_event.wait()
            await asyncio.sleep(0.05)
            t.cancel()
            with pytest.raises(asyncio.CancelledError):
                await t

            # make sure that our previous request, cancelled while waiting for
            # a repsponse, didn't leave the connection open andin a bad state
            with dp.override():
                assert await r.get("bar") == b"bar"
                assert await r.ping()
                assert await r.get("foo") == b"foo"

    await dp.stop()


@pytest.mark.onlynoncluster
@pytest.mark.parametrize("delay", argvalues=[0.05, 0.5, 1, 2])
async def test_standalone_pipeline(delay, redis_addr):
    dp = DelayProxy(addr=("127.0.0.1", 5380), redis_addr=redis_addr, delay=delay * 2)
    await dp.start()
    for b in [True, False]:
        async with Redis(host="127.0.0.1", port=5380, single_connection_client=b) as r:
            with dp.override():
                await r.set("foo", "foo")
                await r.set("bar", "bar")

            pipe = r.pipeline()

            pipe2 = r.pipeline()
            pipe2.get("bar")
            pipe2.ping()
            pipe2.get("foo")

            dp.send_event.clear()
            t = asyncio.create_task(pipe.get("foo").execute())
            # wait until task has settled on the read
            await dp.send_event.wait()
            await asyncio.sleep(0.05)
            t.cancel()
            with pytest.raises(asyncio.CancelledError):
                await t

            # we have now cancelled the pieline in the middle of a request, make sure
            # that the connection is still usable
            with dp.override():
                pipe.get("bar")
                pipe.ping()
                pipe.get("foo")
                await pipe.reset()

                # check that the pipeline is empty after reset
                assert await pipe.execute() == []

                # validating that the pipeline can be used as it could previously
                pipe.get("bar")
                pipe.ping()
                pipe.get("foo")
                assert await pipe.execute() == [b"bar", True, b"foo"]
                assert await pipe2.execute() == [b"bar", True, b"foo"]

    await dp.stop()


@pytest.mark.onlycluster
async def test_cluster(request, redis_addr):

    delay = 0.1
    cluster_port = 6372
    remap_base = 7372
    n_nodes = 6

    remap = []
    proxies = []
    for i in range(n_nodes):
        port = cluster_port + i
        remapped = remap_base + i
        remap.append({"from_port": port, "to_port": remapped})
        forward_addr = redis_addr[0], port
        proxy = DelayProxy(
            addr=("127.0.0.1", remapped), redis_addr=forward_addr, delay=delay
        )
        proxies.append(proxy)

    # start proxies
    await asyncio.gather(*[p.start() for p in proxies])

    def all_clear():
        for p in proxies:
            p.send_event.clear()

    async def wait_for_send():
        asyncio.wait(
            [p.send_event.wait() for p in proxies], return_when=asyncio.FIRST_COMPLETED
        )

    @contextlib.contextmanager
    def override():
        with contextlib.ExitStack() as stack:
            for p in proxies:
                stack.enter_context(p.override())
            yield

    with override():
        r = RedisCluster.from_url(
            f"redis://127.0.0.1:{remap_base}", host_port_remap=remap
        )
        await r.initialize()
        await r.set("foo", "foo")
        await r.set("bar", "bar")

    all_clear()
    t = asyncio.create_task(r.get("foo"))
    # cannot wait on the send event, we don't know which node will be used
    await wait_for_send()
    await asyncio.sleep(delay)
    t.cancel()
    with pytest.raises(asyncio.CancelledError):
        await t

    with override():
        # try a number of requests to excercise all the connections
        async def doit():
            assert await r.get("bar") == b"bar"
            assert await r.ping()
            assert await r.get("foo") == b"foo"

        await asyncio.gather(*[doit() for _ in range(10)])

    await asyncio.gather(*(p.stop() for p in proxies))