add websocket test cases (#11)

* add websocket test cases
* add websockets to python requirements

Author: icing

test/modules/tls/env.py

@@ -107,18 +107,6 @@ def __init__(self, pytestconfig=None):
             '    AddHandler cgi-script .py',
             '    Options +ExecCGI',
             '</Directory>',
-            f'<VirtualHost *:{self.http_port}>',
-            '    ServerName localhost',
-            '    DocumentRoot "htdocs"',
-            '</VirtualHost>',
-            f'<VirtualHost *:{self.http_port}>',
-            f'    ServerName {self.domain_a}',
-            '    DocumentRoot "htdocs/a.mod-tls.test"',
-            '</VirtualHost>',
-            f'<VirtualHost *:{self.http_port}>',
-            f'    ServerName {self.domain_b}',
-            '    DocumentRoot "htdocs/b.mod-tls.test"',
-            '</VirtualHost>',
         ])
         self.add_cert_specs([
             CertificateSpec(domains=[self.domain_a]),

test/modules/tls/test_13_proxy.py

@@ -18,6 +18,8 @@ def _class_scope(self, env):
             ]
         })
         # add vhosts a+b and a ssl proxy from a to b
+        conf.add_vhost('localhost', port=env.http_port)
+        conf.add_vhost(env.domain_b, port=env.http_port, doc_root=f"htdocs/{env.domain_b}")
         conf.add_tls_vhosts(domains=[env.domain_a, env.domain_b])
         conf.install()
         assert env.apache_restart() == 0

test/modules/tls/test_18_ws.py

@@ -0,0 +1,141 @@
+import ssl
+from datetime import datetime, timedelta
+import inspect
+import os
+import shutil
+import subprocess
+import time
+
+import pytest
+import websockets
+from websockets.sync.client import connect
+
+from .conf import TlsTestConf
+
+
+def mk_text_file(fpath: str, lines: int):
+    t110 = 11 * "0123456789"
+    with open(fpath, "w") as fd:
+        for i in range(lines):
+            fd.write("{0:015d}: ".format(i))  # total 128 bytes per line
+            fd.write(t110)
+            fd.write("\n")
+
+
+class TestWebSockets:
+
+    @pytest.fixture(autouse=True, scope='class')
+    def _class_scope(self, env):
+        # Apache config that CONNECT proxies a WebSocket server for paths starting
+        # with '/ws/'
+        # The WebSocket server is started in pytest fixture 'ws_server' below.
+        conf = TlsTestConf(env, extras={
+            'base': [
+                'Timeout 1',
+            ],
+            'localhost': [
+                f'ProxyPass /ws/ http://127.0.0.1:{env.ws_port}/ upgrade=websocket \\',
+                f'timeout=2 flushpackets=on',
+            ],
+            f'cgi.{env.http_tld}': [
+              f'  ProxyPass /ws/ http://127.0.0.1:{env.ws_port}/ \\',
+              f'           upgrade=websocket timeout=2 flushpackets=on',
+              f'  ReadBufferSize 65535'
+            ]
+        })
+        conf.add_vhost('localhost', port=env.http_port)
+        conf.add_tls_vhosts(['localhost'], port=env.https_port)
+        conf.install()
+        mk_text_file(os.path.join(env.gen_dir, "1k.txt"), 8)
+        mk_text_file(os.path.join(env.gen_dir, "10k.txt"), 80)
+        mk_text_file(os.path.join(env.gen_dir, "100k.txt"), 800)
+        mk_text_file(os.path.join(env.gen_dir, "1m.txt"), 8000)
+        mk_text_file(os.path.join(env.gen_dir, "10m.txt"), 80000)
+        assert env.apache_restart() == 0
+
+    def ws_check_alive(self, env, timeout=5):
+        url = f'http://localhost:{env.ws_port}/'
+        end = datetime.now() + timedelta(seconds=timeout)
+        while datetime.now() < end:
+            r = env.curl_get(url, 5)
+            if r.exit_code == 0:
+                return True
+            time.sleep(.1)
+        return False
+
+    def _mkpath(self, path):
+        if not os.path.exists(path):
+            return os.makedirs(path)
+
+    def _rmrf(self, path):
+        if os.path.exists(path):
+            return shutil.rmtree(path)
+
+    def ws_recv_text(self, ws):
+        msg = ""
+        while True:
+            try:
+                msg += ws.recv()
+            except websockets.exceptions.ConnectionClosedOK:
+                return msg
+
+    def ws_recv_bytes(self, ws):
+        msg = b''
+        while True:
+            try:
+                msg += ws.recv()
+            except websockets.exceptions.ConnectionClosedOK:
+                return msg
+
+    @pytest.fixture(autouse=True, scope='class')
+    def ws_server(self, env):
+        # Run our python websockets server that has some special behaviour
+        # for the different path to CONNECT to.
+        run_dir = os.path.join(env.gen_dir, 'ws-server')
+        err_file = os.path.join(run_dir, 'stderr')
+        self._rmrf(run_dir)
+        self._mkpath(run_dir)
+        with open(err_file, 'w') as cerr:
+            cmd = os.path.join(os.path.dirname(inspect.getfile(TestWebSockets)),
+                               'ws_server.py')
+            args = ['python3', cmd, '--port', str(env.ws_port)]
+            p = subprocess.Popen(args=args, cwd=run_dir, stderr=cerr,
+                                 stdout=cerr)
+            if not self.ws_check_alive(env):
+                p.kill()
+                p.wait()
+                pytest.fail(f'ws_server did not start. stderr={open(err_file).readlines()}')
+            yield
+            p.terminate()
+
+    def test_tls_18_01_direct(self, env):
+        with connect(f"ws://127.0.0.1:{env.ws_port}/echo") as ws:
+            message = "Hello world!"
+            ws.send(message)
+            response = self.ws_recv_text(ws)
+            assert response == message
+
+    def test_tls_18_02_httpd_plain(self, env):
+        with connect(f"ws://localhost:{env.http_port}/ws/echo/") as ws:
+            message = "Hello world!"
+            ws.send(message)
+            response = self.ws_recv_text(ws)
+            assert response == message
+
+    @pytest.mark.parametrize("fname", ["1k.txt", "10k.txt", "100k.txt", "1m.txt", "10m.txt"])
+    def test_tls_18_03_file(self, env, fname):
+        expected = open(os.path.join(env.gen_dir, fname), 'rb').read()
+        with connect(f"ws://localhost:{env.http_port}/ws/file/{fname}") as ws:
+            response = self.ws_recv_bytes(ws)
+            assert response == expected
+
+    @pytest.mark.parametrize("fname", ["1k.txt", "10k.txt", "100k.txt", "1m.txt", "10m.txt"])
+    def test_tls_18_04_tls_file(self, env, fname):
+        expected = open(os.path.join(env.gen_dir, fname), 'rb').read()
+        ssl_ctx = ssl.SSLContext(protocol=ssl.PROTOCOL_TLS_CLIENT)
+        ssl_ctx.check_hostname = False
+        ssl_ctx.verify_mode = ssl.VerifyMode.CERT_NONE
+        with connect(f"wss://localhost:{env.https_port}/ws/file/{fname}",
+                     ssl_context=ssl_ctx) as ws:
+            response = self.ws_recv_bytes(ws)
+            assert response == expected

test/modules/tls/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())

test/pyhttpd/env.py

@@ -312,6 +312,7 @@ def __init__(self, pytestconfig=None):
             f"test2.{self._http_tld}",
             f"test3.{self._http_tld}",
             f"cgi.{self._http_tld}",
+            "localhost",
         ], key_type='rsa4096')]
 
         self._verify_certs = False

test/requirements.txt

@@ -19,3 +19,4 @@ filelock
 python-multipart
 psutil
 tqdm
+websockets