proxy.pyx

import asyncio
import errno
import signal
import socket

import uvloop


cdef str HOST = "127.0.0.1"
cdef int PORT = 8888
cdef int TIMEOUT = 10
cdef int SOCK_BUF_SIZE = 131072
cdef int INITIAL_READ_SIZE = 8192
cdef int MAX_HEADER_BYTES = 65536
cdef unsigned long BW_FLUSH_BYTES = 1 << 20
cdef bytes AUTH_VALUE = b"Basic dXNlcm5hbWU6cGFzc3dvcmQ="
cdef bytes RESP_200 = b"HTTP/1.1 200 Connection Established\r\n\r\n"
cdef bytes RESP_407 = (
    b"HTTP/1.1 407 Proxy Auth Required\r\n"
    b"Proxy-Authenticate: Basic\r\n\r\n"
)
cdef bytes RESP_431 = b"HTTP/1.1 431 Request Header Fields Too Large\r\n\r\n"
cdef bytes RESP_502 = b"HTTP/1.1 502 Bad Gateway\r\n\r\n"

cdef unsigned long bw = 0
cdef unsigned long total_requests = 0
cdef unsigned long successful_requests = 0
cdef unsigned long failed_requests = 0
cdef unsigned long client_disconnects = 0

cdef int RELAY_FAILED = -1
cdef int RELAY_CLIENT_DISCONNECTED = 0
cdef int RELAY_COMPLETE = 1


cdef str format_bw(unsigned long transferred):
    if transferred < 1024:
        return f"{transferred}B"
    if transferred < 1024**2:
        return f"{transferred / 1024:.2f}KB"
    return f"{transferred / 1024**2:.2f}MB"


cdef inline void tune_socket(object sock):
    sock.setblocking(False)
    try:
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    except OSError:
        pass


def print_metrics():
    """Print the counters captured by benchmark result logs."""
    print("Metrics:")
    print(f"Bandwidth usage: {format_bw(bw)}")
    print(f"Total Requests: {total_requests}")
    print(f"Successful Visits: {successful_requests}")
    print(f"Failed Visits: {failed_requests}")
    print(f"Client Disconnects: {client_disconnects}")


cdef inline bint valid_auth(bytes data, Py_ssize_t header_end):
    """Accept exactly one correctly formed Proxy-Authorization header."""
    cdef Py_ssize_t request_end = data.find(b"\r\n", 0, header_end)
    cdef Py_ssize_t line_start
    cdef Py_ssize_t line_end
    cdef Py_ssize_t colon
    cdef int auth_count = 0
    cdef bytes name
    cdef bytes value

    if request_end == -1:
        return False
    line_start = request_end + 2
    while line_start < header_end:
        line_end = data.find(b"\r\n", line_start, header_end + 2)
        if line_end == -1:
            return False
        colon = data.find(b":", line_start, line_end)
        if colon <= line_start:
            return False
        name = data[line_start:colon]
        if b" " in name or b"\t" in name:
            return False
        if name.lower() == b"proxy-authorization":
            auth_count += 1
            value = data[colon + 1:line_end].strip(b" \t")
            if value != AUTH_VALUE:
                return False
        line_start = line_end + 2

    return auth_count == 1


cdef inline int relay_error_status(object error, bint client_endpoint):
    cdef object error_number = getattr(error, "errno", None)
    if client_endpoint and (
        isinstance(
            error,
            (ConnectionResetError, ConnectionAbortedError, BrokenPipeError),
        )
        or error_number in (errno.ECONNRESET, errno.ECONNABORTED, errno.EPIPE)
    ):
        return RELAY_CLIENT_DISCONNECTED
    return RELAY_FAILED


async def relay_sock(
    object loop,
    object source,
    object destination,
    bint source_is_client,
):
    """Relay one direction and propagate a clean EOF as a half-close."""
    global bw
    cdef bytes data
    cdef unsigned long transferred = 0
    cdef bint client_endpoint = False

    try:
        while True:
            client_endpoint = source_is_client
            data = await loop.sock_recv(source, SOCK_BUF_SIZE)
            if not data:
                break
            transferred += len(data)
            if transferred >= BW_FLUSH_BYTES:
                bw += transferred
                transferred = 0
            client_endpoint = not source_is_client
            await loop.sock_sendall(destination, data)
    except OSError as error:
        return relay_error_status(error, client_endpoint)
    finally:
        bw += transferred

    try:
        destination.shutdown(socket.SHUT_WR)
    except OSError:
        pass
    return RELAY_COMPLETE


async def relay_pair(object loop, object client_sock, object remote_sock):
    """Relay both directions until both cleanly half-close or either fails."""
    to_client = loop.create_task(
        relay_sock(loop, remote_sock, client_sock, False)
    )
    to_remote = loop.create_task(
        relay_sock(loop, client_sock, remote_sock, True)
    )
    relay_tasks = (to_client, to_remote)
    try:
        done, pending = await asyncio.wait(
            relay_tasks, return_when=asyncio.FIRST_COMPLETED
        )
        results = [task.result() for task in done]
        if pending and any(result != RELAY_COMPLETE for result in results):
            for task in pending:
                task.cancel()
            trailing = await asyncio.gather(*pending, return_exceptions=True)
            results.extend(
                result
                for result in trailing
                if not isinstance(result, BaseException)
            )
        elif pending:
            results.extend(await asyncio.gather(*pending))
        if RELAY_FAILED in results:
            return RELAY_FAILED
        if RELAY_CLIENT_DISCONNECTED in results:
            return RELAY_CLIENT_DISCONNECTED
        return RELAY_COMPLETE
    except BaseException:
        for task in relay_tasks:
            if not task.done():
                task.cancel()
        await asyncio.gather(*relay_tasks, return_exceptions=True)
        raise


async def handle_client(object client_sock):
    """Authenticate one proxy connection and relay it to its destination."""
    global total_requests, successful_requests, failed_requests
    global client_disconnects
    cdef Py_ssize_t header_end
    cdef Py_ssize_t target_end
    cdef bytes data
    cdef bytes chunk
    cdef bytes target
    cdef bytes connect_host
    cdef bytes port_bytes
    cdef bytes initial_tunnel_data
    cdef int port
    cdef int relay_result

    total_requests += 1
    remote_sock = None
    try:
        loop = asyncio.get_running_loop()
        tune_socket(client_sock)

        async with asyncio.timeout_at(loop.time() + TIMEOUT):
            data = await loop.sock_recv(client_sock, INITIAL_READ_SIZE)
            if not data:
                return
            header_end = data.find(b"\r\n\r\n")
            if header_end == -1:
                request_buf = bytearray(data)
                while header_end == -1:
                    if len(request_buf) >= MAX_HEADER_BYTES:
                        await loop.sock_sendall(client_sock, RESP_431)
                        failed_requests += 1
                        return
                    chunk = await loop.sock_recv(
                        client_sock,
                        min(INITIAL_READ_SIZE, MAX_HEADER_BYTES - len(request_buf)),
                    )
                    if not chunk:
                        return
                    request_buf.extend(chunk)
                    header_end = request_buf.find(b"\r\n\r\n")
                data = bytes(request_buf)

        if not valid_auth(data, header_end):
            await loop.sock_sendall(client_sock, RESP_407)
            return

        if not data.startswith(b"CONNECT "):
            raise ValueError("CONNECT required")
        target_end = data.find(b" ", 8, header_end)
        if target_end == -1:
            raise ValueError("malformed CONNECT request line")
        target = data[8:target_end]
        try:
            connect_host, port_bytes = target.rsplit(b":", 1)
            if not connect_host:
                raise ValueError("empty CONNECT host")
            port = int(port_bytes)
            if port < 1 or port > 65535:
                raise ValueError("invalid CONNECT port")
            remote_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            tune_socket(remote_sock)
            await loop.sock_connect(remote_sock, (connect_host.decode("ascii"), port))
        except Exception as error:
            print(f"Error connecting to {target!r}: {error}")
            failed_requests += 1
            try:
                await loop.sock_sendall(client_sock, RESP_502)
            except OSError:
                pass
            return

        try:
            await loop.sock_sendall(client_sock, RESP_200)
        except OSError as error:
            if relay_error_status(error, True) == RELAY_CLIENT_DISCONNECTED:
                client_disconnects += 1
            else:
                failed_requests += 1
            return

        try:
            initial_tunnel_data = data[header_end + 4:]
            if initial_tunnel_data:
                await loop.sock_sendall(remote_sock, initial_tunnel_data)
        except OSError:
            failed_requests += 1
            return

        relay_result = await relay_pair(loop, client_sock, remote_sock)
        if relay_result == RELAY_FAILED:
            failed_requests += 1
        elif relay_result == RELAY_CLIENT_DISCONNECTED:
            client_disconnects += 1
        else:
            successful_requests += 1
    except asyncio.TimeoutError:
        failed_requests += 1
    except Exception as error:
        print(f"Error: {error}")
        failed_requests += 1
    finally:
        client_sock.close()
        if remote_sock is not None:
            remote_sock.close()


async def main():
    """Run one uvloop worker with a raw socket accept loop."""
    loop = asyncio.get_running_loop()
    server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    try:
        server_sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_FASTOPEN, 1024)
    except OSError:
        pass
    server_sock.setblocking(False)
    server_sock.bind((HOST, PORT))
    server_sock.listen(65535)
    print(f"Running on {HOST}:{PORT}")

    stop_event = asyncio.Event()
    main_task = asyncio.current_task()

    def request_stop():
        if stop_event.is_set():
            return
        print_metrics()
        stop_event.set()
        main_task.cancel()

    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, request_stop)

    try:
        while not stop_event.is_set():
            try:
                client_sock, _ = await loop.sock_accept(server_sock)
            except (OSError, asyncio.CancelledError):
                if stop_event.is_set():
                    break
                raise
            loop.create_task(handle_client(client_sock))
    finally:
        server_sock.close()
        client_tasks = tuple(
            task for task in asyncio.all_tasks() if task is not main_task
        )
        for client_task in client_tasks:
            client_task.cancel()
        if client_tasks:
            await asyncio.gather(*client_tasks, return_exceptions=True)


uvloop.run(main())