Properly close sockets for the cwe_404 proxy

Author: kristjanvalur

tests/test_asyncio/test_cwe_404.py

@@ -49,22 +49,22 @@ def set_delay(self, delay: float = 0.0):
     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()
+        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)
 
     async def stop(self):
-        # clean up enough so that we can reuse the looper
+        # shutdown the server
         self.task.cancel()
         try:
             await self.task
         except asyncio.CancelledError:
             pass
+        await self.server.wait_closed()
+        # do we need to close individual connections too?
+        # prudently close all async generators
         loop = self.server.get_loop()
         await loop.shutdown_asyncgens()
 
@@ -75,16 +75,20 @@ async def pipe(
         name="",
         event: asyncio.Event = None,
     ):
-        while True:
-            data = await reader.read(1000)
-            if not data:
-                break
-            # print(f"{name} read {len(data)} delay {self.delay}")
-            if event:
-                event.set()
-            await asyncio.sleep(self.delay)
-            writer.write(data)
-            await writer.drain()
+        try:
+            while True:
+                data = await reader.read(1000)
+                if not data:
+                    break
+                # print(f"{name} read {len(data)} delay {self.delay}")
+                if event:
+                    event.set()
+                await asyncio.sleep(self.delay)
+                writer.write(data)
+                await writer.drain()
+        finally:
+            writer.close()
+            await writer.wait_closed()
 
 
 @pytest.mark.onlynoncluster