1 files changed, 104 insertions, 0 deletions
diff --git a/test/modules/http2/ws_server.py b/test/modules/http2/ws_server.py
new file mode 100644
index 0000000..99fb9cf
--- /dev/null
+++ b/test/modules/http2/ws_server.py
@@ -0,0 +1,104 @@
+#!/usr/bin/env python3
+import argparse
+import asyncio
+import logging
+import os
+import sys
+import time
+
+import websockets.server as ws_server
+from websockets.exceptions import ConnectionClosedError
+
+log = logging.getLogger(__name__)
+
+logging.basicConfig(
+    format="[%(asctime)s] %(message)s",
+    level=logging.DEBUG,
+)
+
+
+async def echo(websocket):
+    try:
+        async for message in websocket:
+            try:
+                log.info(f'got request {message}')
+            except Exception as e:
+                log.error(f'error {e} getting path from {message}')
+            await websocket.send(message)
+    except ConnectionClosedError:
+        pass
+
+
+async def on_async_conn(conn):
+    rpath = str(conn.path)
+    pcomps = rpath[1:].split('/')
+    if len(pcomps) == 0:
+        pcomps = ['echo']  # default handler
+    log.info(f'connection for {pcomps}')
+    if pcomps[0] == 'echo':
+        log.info(f'/echo endpoint')
+        for message in await conn.recv():
+            await conn.send(message)
+    elif pcomps[0] == 'text':
+        await conn.send('hello!')
+    elif pcomps[0] == 'file':
+        if len(pcomps) < 2:
+            conn.close(code=4999, reason='unknown file')
+            return
+        fpath = os.path.join('../', pcomps[1])
+        if not os.path.exists(fpath):
+            conn.close(code=4999, reason='file not found')
+            return
+        bufsize = 0
+        if len(pcomps) > 2:
+            bufsize = int(pcomps[2])
+        if bufsize <= 0:
+            bufsize = 16*1024
+        delay_ms = 0
+        if len(pcomps) > 3:
+            delay_ms = int(pcomps[3])
+        n = 1
+        if len(pcomps) > 4:
+            n = int(pcomps[4])
+        for _ in range(n):
+            with open(fpath, 'r+b') as fd:
+                while True:
+                    buf = fd.read(bufsize)
+                    if buf is None or len(buf) == 0:
+                        break
+                    await conn.send(buf)
+                    if delay_ms > 0:
+                        time.sleep(delay_ms/1000)
+    else:
+        log.info(f'unknown endpoint: {rpath}')
+        await conn.close(code=4999, reason='path unknown')
+    await conn.close(code=1000, reason='')
+
+
+async def run_server(port):
+    log.info(f'starting server on port {port}')
+    async with ws_server.serve(ws_handler=on_async_conn,
+                               host="localhost", port=port):
+        await asyncio.Future()
+
+
+async def main():
+    parser = argparse.ArgumentParser(prog='scorecard',
+                                     description="Run a websocket echo server.")
+    parser.add_argument("--port", type=int,
+                        default=0, help="port to listen on")
+    args = parser.parse_args()
+
+    if args.port == 0:
+        sys.stderr.write('need --port\n')
+        sys.exit(1)
+
+    logging.basicConfig(
+        format="%(asctime)s %(message)s",
+        level=logging.DEBUG,
+    )
+    await run_server(args.port)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())