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(":")
    port = port.split("/")[0]
    return host, int(port)


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()

    async def start(self):
        # test that we can connect to redis
        async 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 set_delay(self, delay: float = 0.0):
        """
        Allow to override the delay for parts of tests which aren't time dependent,
        to speed up execution.
        """
        old_delay = self.delay
        self.delay = delay
        try:
            yield
        finally:
            self.delay = old_delay

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

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

    async def stop(self):
        # clean up enough so that we can reuse the looper
        self.ROUTINE.cancel()
        try:
            await self.ROUTINE
        except asyncio.CancelledError:
            pass
        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.set_delay(0):
                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.set_delay(0):
                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.set_delay(0):
                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.set_delay(0):
                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

    def remap(host, port):
        return host, remap_base + port - cluster_port

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

    # helpers to work with all or any proxy
    def all_clear():
        for p in proxies:
            p.send_event.clear()

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

    @contextlib.contextmanager
    def set_delay(delay: int = 0):
        with contextlib.ExitStack() as stack:
            for p in proxies:
                stack.enter_context(p.delay_as(delay))
            yield

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

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

        all_clear()
        with set_delay(delay=delay):
            t = asyncio.create_task(r.get("foo"))
            # One of the proxies will handle our request, wait for it to send
            await any_wait()
            await asyncio.sleep(delay)
            t.cancel()
            with pytest.raises(asyncio.CancelledError):
                await t

        # 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)])

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