#!/usr/bin/env python3
"""Streaming HTTP proxy integration tests: one worker, real sockets and TLS."""
import asyncio
import concurrent.futures
import contextlib
import hashlib
import http.client
import json
import os
import pathlib
import re
import socket
import ssl
import subprocess
import tempfile
import threading
import time
import unittest
from test_websocket import Backend as WSBackend, frame, recv_frame, exact

ROOT = pathlib.Path(__file__).resolve().parents[1]
BINARY = pathlib.Path(os.environ.get('SUPERWEB', ROOT / 'superweb')).resolve()
BLOCK = bytes(range(256)) * 256


class Backend:
    def __init__(self):
        self.loop = asyncio.new_event_loop()
        self.ready = threading.Event()
        self.seen = threading.Event()
        self.upload_seen = threading.Event()
        self.release = threading.Event()
        self.thread = threading.Thread(target=self.main, daemon=True)
        self.thread.start()
        assert self.ready.wait(5)

    def main(self):
        asyncio.set_event_loop(self.loop)
        self.server = self.loop.run_until_complete(asyncio.start_server(self.handle, '127.0.0.1', 0))
        self.port = self.server.sockets[0].getsockname()[1]
        self.ready.set()
        self.loop.run_forever()
        self.server.close()
        self.loop.run_until_complete(self.server.wait_closed())
        tasks = asyncio.all_tasks(self.loop)
        for task in tasks:
            task.cancel()
        self.loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))
        self.loop.close()

    def close(self):
        self.loop.call_soon_threadsafe(self.loop.stop)
        self.thread.join(5)

    async def handle(self, reader, writer):
        try:
            header = await reader.readuntil(b'\r\n\r\n')
            method, target, _ = header.split(b'\r\n', 1)[0].split(b' ')
            path = target.split(b'?', 1)[0]
            headers = dict(line.lower().split(b':', 1) for line in header.split(b'\r\n')[1:] if b':' in line)
            self.seen.set()
            if path == b'/delay':
                await asyncio.sleep(.8)
            if path == b'/stall':
                await asyncio.sleep(10)
                return
            if path == b'/reject':
                writer.write(b'HTTP/1.1 413 Content Too Large\r\nContent-Length: 0\r\n\r\n')
                await writer.drain()
                await asyncio.sleep(.1)
                return
            if path in (b'/head', b'/empty', b'/modified'):
                writer.write({b'/head': b'HTTP/1.1 200 OK\r\nContent-Length: 1000000\r\n\r\n',
                              b'/empty': b'HTTP/1.1 204 No Content\r\n\r\n',
                              b'/modified': b'HTTP/1.1 304 Not Modified\r\nContent-Length: 123\r\n\r\n'}[path])
                await writer.drain()
                await asyncio.sleep(2)
                return
            if path == b'/trickle':
                for value in b'HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok':
                    writer.write(bytes([value])); await writer.drain(); await asyncio.sleep(.08)
                return
            if path == b'/truncated':
                writer.write(b'HTTP/1.1 200 OK\r\nContent-Length: 20\r\n\r\nshort')
                await writer.drain()
                await asyncio.sleep(.1)
                return
            if path.startswith(b'/bad/'):
                examples = [b'Content-Length: 2\r\nContent-Length: 3',
                            b'Content-Length: 2\r\nTransfer-Encoding: chunked',
                            b'Transfer-Encoding: chunked, gzip', b'Bad Header: value',
                            b'X-Unsafe: value\x00tail', b'Content-Length: -1',
                            b' Transfer-Encoding: chunked', b'Transfer-Encoding: notchunked']
                writer.write(b'HTTP/1.1 200 OK\r\n' + examples[int(path.rsplit(b'/', 1)[1])] + b'\r\n\r\nxx')
                return
            if path in (b'/chunks', b'/stream', b'/badchunk'):
                writer.write(b'HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nTrailer: X-End\r\n\r\n')
                await writer.drain()
                if path == b'/badchunk':
                    writer.write(b'4\r\ntestXX0\r\n\r\n'); return
                if path == b'/stream':
                    writer.write(b'5\r\nfirst\r\n'); await writer.drain()
                    while not self.release.is_set(): await asyncio.sleep(.005)
                    writer.write(b'4\r\nlast\r\n0\r\n\r\n'); return
                for value in b'4;tag=test\r\nwiki\r\n5\r\npedia\r\n0\r\nX-End: yes\r\n\r\n':
                    writer.write(bytes([value])); await writer.drain(); await asyncio.sleep(.001)
                return
            if path == b'/eof':
                writer.write(b'HTTP/1.0 200 OK\r\n\r\nclose delimited'); return
            if path == b'/large':
                writer.write(b'HTTP/1.1 200 OK\r\nContent-Length: 16777216\r\n\r\n')
                for _ in range(256):
                    writer.write(BLOCK); await writer.drain()
                return
            if path == b'/slowread':
                await asyncio.sleep(1)
            if headers.get(b'transfer-encoding', b'').strip() == b'chunked':
                chunks = bytearray()
                while True:
                    size = int((await reader.readline()).strip(), 16)
                    if not size:
                        assert await reader.readexactly(2) == b'\r\n'
                        break
                    self.upload_seen.set()
                    chunks.extend(await reader.readexactly(size))
                    assert await reader.readexactly(2) == b'\r\n'
                reader = asyncio.StreamReader(); reader.feed_data(chunks); reader.feed_eof()
                length = len(chunks)
            else:
                length = int(headers.get(b'content-length', b'0'))
            digest = hashlib.sha256()
            left = length
            body = b''
            while left:
                data = await reader.read(min(left, 4096))
                if not data: return
                self.upload_seen.set()
                digest.update(data); left -= len(data)
                if path == b'/echo': body += data
            if path != b'/echo':
                body = json.dumps({'path': target.decode(), 'bytes': length, 'sha256': digest.hexdigest(),
                                   'headers': {k.decode(): v.strip().decode() for k, v in headers.items()}}).encode()
            if path == b'/hints':
                writer.write(b'HTTP/1.1 100 Continue\r\n\r\nHTTP/1.1 103 Early Hints\r\nLink: </x>\r\n\r\n')
            writer.write(b'HTTP/1.1 200 OK\r\nContent-Length: ' + str(len(body)).encode() +
                         b'\r\nConnection: X-Remove, close\r\nX-Remove: secret\r\nSet-Cookie: a=1\r\nSet-Cookie: b=2\r\n\r\n' + body)
            await writer.drain()
        except (asyncio.IncompleteReadError, ConnectionError, OSError):
            pass
        finally:
            writer.close()
            with contextlib.suppress(ConnectionError, OSError): await writer.wait_closed()


class ProxyTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.temp = tempfile.TemporaryDirectory(prefix='sw-http-proxy-')
        cls.root = pathlib.Path(cls.temp.name)
        (cls.root / 'hello.txt').write_bytes(b'assets stay responsive')
        cls.key, cls.cert = cls.root / 'key.pem', cls.root / 'cert.pem'
        subprocess.run(['openssl', 'req', '-x509', '-newkey', 'rsa:2048', '-nodes', '-days', '1', '-keyout', str(cls.key), '-out', str(cls.cert), '-subj', '/CN=localhost', '-addext', 'subjectAltName=DNS:localhost,IP:127.0.0.1'], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        cls.tls = 'x64' not in BINARY.name
        cls.ctx = ssl.create_default_context(cafile=str(cls.cert))
        cls.ctx.set_alpn_protocols(['http/1.1'])
        cls.ws_backend = WSBackend()

    @classmethod
    def tearDownClass(cls):
        cls.ws_backend.close()
        cls.temp.cleanup()

    def setUp(self):
        self.backend = Backend()
        self.sockets = []
        self.processes = []

    def tearDown(self):
        for sock in self.sockets: sock.close()
        for proc in self.processes:
            if proc.poll() is None:
                proc.terminate()
                try: proc.wait(4)
                except subprocess.TimeoutExpired: proc.kill(); proc.wait()
            errors = proc.stderr.read()
            self.assertNotIn('ERROR: AddressSanitizer', errors)
            self.assertNotIn('runtime error:', errors)
            self.assertEqual(proc.returncode, 0, errors)
            peaks = re.findall(r'proxy_peak_queued=(\d+)', errors)
            limit = 4096 if '--proxy-buffer' in proc.args else 16384
            for peak in peaks: self.assertLessEqual(int(peak), getattr(self, 'h2_peak_limit', 2 * limit))
            proc.stderr.close(); proc.stdout.close()
        self.backend.close()

    def start(self, tls=False, extra=()):
        command = [str(BINARY), '--root', str(self.root), '--bind', '127.0.0.1', '--port', '0', '--workers', '1',
                   '--proxy', f'/api=localhost:{self.backend.port}', '--proxy', f'/game=127.0.0.1:{self.ws_backend.port}',
                   '--proxy', '/down=127.0.0.1:1', '--drain-ms', '100', *extra]
        if tls: command += ['--tls-port', '0', '--tls-cert', str(self.cert), '--tls-key', str(self.key)]
        proc = subprocess.Popen(command, env={**os.environ, 'SW_WORKER_STATS': '1'}, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        self.processes.append(proc)
        line = proc.stdout.readline()
        self.assertIn('superweb listening', line)
        return int(re.search(r'https://[^:]+:(\d+)' if tls else r'http://[^:]+:(\d+)', line)[1])

    def connect(self, port, tls=False, small=False):
        sock = socket.socket(); sock.settimeout(3)
        if small: sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 4096)
        sock.connect(('127.0.0.1', port))
        if tls: sock = self.ctx.wrap_socket(sock, server_hostname='localhost')
        self.sockets.append(sock)
        return sock

    def get(self, port, path='/hello.txt', tls=False):
        sock = self.connect(port, tls)
        sock.sendall(f'GET {path} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n'.encode())
        response = http.client.HTTPResponse(sock); response.begin()
        return response.status, response.read()

    def websocket(self, port, tls=False):
        sock = self.connect(port, tls)
        sock.sendall(b'GET /game/echo HTTP/1.1\r\nHost: localhost\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\r\n')
        header = b''
        while not header.endswith(b'\r\n\r\n'): header += exact(sock, 1)
        self.assertTrue(header.startswith(b'HTTP/1.1 101'))
        return sock

    def responsive(self, port, ws, tls=False):
        start = time.monotonic()
        ws.sendall(frame(b'game tick'))
        self.assertEqual(recv_frame(ws), (130, b'game tick'))
        self.assertEqual(self.get(port, tls=tls), (200, b'assets stay responsive'))
        self.assertLess(time.monotonic() - start, .5)

    def test_slow_backends_do_not_block_websockets_or_assets(self):
        port = self.start()
        ws = self.websocket(port)
        sessions = [self.connect(port) for _ in range(32)]
        start = time.monotonic()
        for sock in sessions: sock.sendall(b'GET /api/delay HTTP/1.1\r\nHost: localhost\r\n\r\n')
        self.assertTrue(self.backend.seen.wait(2))
        self.responsive(port, ws)
        for sock in sessions:
            response = http.client.HTTPResponse(sock); response.begin()
            self.assertEqual(response.status, 200); response.read()
        self.assertLess(time.monotonic() - start, 2)

    def test_stream_upload_before_client_finishes_and_keepalive_pipeline(self):
        port = self.start()
        sock = self.connect(port)
        sock.sendall(b'POST /api/echo HTTP/1.1\r\nHost: localhost\r\nContent-Length: 65536\r\nExpect: 100-continue\r\n\r\n')
        self.assertEqual(exact(sock, 25), b'HTTP/1.1 100 Continue\r\n\r\n')
        sock.sendall(BLOCK[:4096])
        self.assertTrue(self.backend.upload_seen.wait(1), 'upload was buffered instead of streamed')
        ws = self.websocket(port); self.responsive(port, ws)
        sock.sendall(BLOCK[4096:] + b'GET /hello.txt HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n')
        # One buffered file reader preserves bytes from the pipelined response.
        reader = sock.makefile('rb')
        self.assertTrue(reader.readline().startswith(b'HTTP/1.1 200'))
        headers = http.client.parse_headers(reader)
        self.assertEqual(int(headers['Content-Length']), len(BLOCK))
        self.assertEqual(reader.read(len(BLOCK)), BLOCK)
        self.assertTrue(reader.readline().startswith(b'HTTP/1.1 200'))
        headers = http.client.parse_headers(reader)
        self.assertEqual(reader.read(int(headers['Content-Length'])), b'assets stay responsive')
        reader.close()

    def test_stream_response_before_backend_finishes(self):
        port = self.start(); sock = self.connect(port)
        sock.sendall(b'GET /api/stream HTTP/1.1\r\nHost: localhost\r\n\r\n')
        response = http.client.HTTPResponse(sock); response.begin()
        self.assertEqual(response.read(5), b'first')
        self.responsive(port, self.websocket(port))
        self.backend.release.set()
        self.assertEqual(response.read(), b'last')
        sock.sendall(b'GET /api/hints HTTP/1.1\r\nHost: localhost\r\n\r\n')
        response = http.client.HTTPResponse(sock); response.begin()
        self.assertEqual(response.status, 200); response.read()

    def test_framing_head_empty_chunked_eof_and_hop_headers(self):
        port = self.start()
        for path, method, status, length in [('/head', 'HEAD', 200, '1000000'), ('/empty', 'GET', 204, None), ('/modified', 'GET', 304, '123')]:
            sock = self.connect(port); start = time.monotonic()
            sock.sendall(f'{method} /api{path} HTTP/1.1\r\nHost: localhost\r\n\r\n'.encode())
            response = http.client.HTTPResponse(sock, method=method); response.begin()
            self.assertEqual(response.status, status); self.assertEqual(response.read(), b'')
            self.assertEqual(response.getheader('Content-Length'), length)
            self.assertLess(time.monotonic() - start, .5)
        for version in ('1.0', '1.1'):
            for path, expected in [('/chunks', b'wikipedia'), ('/eof', b'close delimited')]:
                sock = self.connect(port)
                sock.sendall(f'GET /api{path} HTTP/{version}\r\nHost: localhost\r\n\r\n'.encode())
                response = http.client.HTTPResponse(sock); response.begin()
                self.assertEqual(response.read(), expected)
                self.assertIsNone(response.getheader('Trailer'))
        sock = self.connect(port)
        sock.sendall(b'GET /api?query=1 HTTP/1.1\r\nHost: localhost\r\nConnection: X-Remove\r\nX-Remove: secret\r\nX-Forwarded-For: forged\r\nX-Keep: yes\r\n\r\n')
        response = http.client.HTTPResponse(sock); response.begin()
        self.assertIsNone(response.getheader('X-Remove'))
        self.assertEqual([v for k,v in response.getheaders() if k.lower() == 'set-cookie'], ['a=1','b=2'])
        body = json.loads(response.read())
        self.assertEqual(body['path'], '/?query=1')
        self.assertNotIn('x-remove', body['headers'])
        self.assertEqual(body['headers']['x-forwarded-for'], '127.0.0.1')
        self.assertEqual(body['headers']['x-keep'], 'yes')

    def test_bad_upstreams_early_rejection_and_deadlines(self):
        port = self.start(extra=['--proxy-idle-ms','300', '--proxy-timeout-ms','600'])
        for i in range(8): self.assertEqual(self.get(port, f'/api/bad/{i}')[0], 502)
        self.assertEqual(self.get(port, '/down/')[0], 502)
        for path in ('/stall', '/trickle'):
            start = time.monotonic()
            self.assertEqual(self.get(port, '/api' + path)[0], 504)
            self.assertLess(time.monotonic() - start, 1.3)
        sock = self.connect(port)
        sock.sendall(b'POST /api/echo HTTP/1.1\r\nHost: localhost\r\nContent-Length: 100\r\n\r\nx')
        response = http.client.HTTPResponse(sock); response.begin()
        self.assertEqual(response.status, 408); response.read()
        sock = self.connect(port)
        sock.sendall(b'POST /api/reject HTTP/1.1\r\nHost: localhost\r\nContent-Length: 1000000\r\n\r\nx')
        response = http.client.HTTPResponse(sock); response.begin()
        self.assertEqual(response.status, 413); self.assertEqual(response.getheader('Connection'), 'close')
        self.assertEqual(response.read(), b'')
        for path in ('/truncated', '/badchunk'):
            sock = self.connect(port)
            sock.sendall(f'GET /api{path} HTTP/1.1\r\nHost: localhost\r\n\r\n'.encode())
            response = http.client.HTTPResponse(sock); response.begin()
            if response.status == 502: response.read()
            else:
                with self.assertRaises(http.client.IncompleteRead): response.read()
        self.assertEqual(self.get(port)[0], 200)

    def test_large_tls_stream_paused_reader_and_upload_backpressure(self):
        if not self.tls: self.skipTest('TLS unavailable in x64 build')
        port = self.start(tls=True, extra=['--proxy-buffer','4096', '--socket-buffer','8192'])
        ws = self.websocket(port, tls=True)
        sock = self.connect(port, tls=True, small=True)
        sock.sendall(b'GET /api/large HTTP/1.1\r\nHost: localhost\r\n\r\n')
        response = http.client.HTTPResponse(sock); response.begin()
        time.sleep(.3)
        self.responsive(port, ws, tls=True)
        digest = hashlib.sha256(); length = 0
        while data := response.read(65536): digest.update(data); length += len(data)
        self.assertEqual(length, 16777216)
        self.assertEqual(digest.digest(), hashlib.sha256(BLOCK * 256).digest())
        upload = self.connect(port, tls=True)
        upload.sendall(b'POST /api/slowread HTTP/1.1\r\nHost: localhost\r\nContent-Length: 1048576\r\n\r\n')
        with concurrent.futures.ThreadPoolExecutor(1) as pool:
            pending = pool.submit(upload.sendall, BLOCK * 16)
            time.sleep(.15); self.responsive(port, ws, tls=True)
            pending.result(5)
        # Linux's deliberately tiny TCP windows can stretch this bulk transfer.
        # Keep the independent game/asset probes at their 500 ms deadline.
        upload.settimeout(20)
        result = http.client.HTTPResponse(upload); result.begin()
        body = json.loads(result.read())
        self.assertEqual(body['bytes'], 1048576)
        self.assertEqual(body['sha256'], hashlib.sha256(BLOCK * 16).hexdigest())
        self.assertEqual(body['headers']['x-forwarded-proto'], 'https')

    def test_tls_stalled_reader_timeout_and_client_cancellation(self):
        if not self.tls: self.skipTest('TLS unavailable in x64 build')
        port = self.start(tls=True, extra=['--proxy-idle-ms','300', '--proxy-buffer','4096', '--socket-buffer','8192'])
        ws = self.websocket(port, tls=True)
        sock = self.connect(port, tls=True, small=True)
        sock.sendall(b'GET /api/large HTTP/1.1\r\nHost: localhost\r\n\r\n')
        response = http.client.HTTPResponse(sock); response.begin()
        time.sleep(1)
        self.responsive(port, ws, tls=True)
        with self.assertRaises(http.client.IncompleteRead): response.read()
        for _ in range(32):
            cancelled = self.connect(port, tls=True)
            cancelled.sendall(b'POST /api/echo HTTP/1.1\r\nHost: localhost\r\nContent-Length: 1000\r\n\r\nx')
            cancelled.close()
        self.responsive(port, ws, tls=True)

    def test_disconnect_churn_and_shutdown(self):
        port = self.start()
        for i in range(100):
            sock = self.connect(port)
            sock.sendall(b'GET /api/delay HTTP/1.1\r\nHost: localhost\r\n\r\n')
            sock.close()
        self.responsive(port, self.websocket(port))
        proc = self.processes[-1]
        start = time.monotonic(); proc.terminate(); proc.wait(2)
        self.assertLess(time.monotonic() - start, 1)


if __name__ == '__main__': unittest.main()