Skip to content

Building a TCP Proxy with asyncio Streams

A TCP proxy accepts a client connection, opens one to an upstream, and copies bytes both ways until both sides are done. It is the core of load balancers, port forwarders, test fault injectors and protocol sniffers, and with asyncio streams it is about thirty lines. The details that make it correct are about ending connections. Tested on Python 3.14 with an upstream that reads until end-of-file and then replies: a proxy that closed both sides when the client finished sending returned an empty response; a proxy that forwarded the half-close with write_eof() returned the upstream's full reply. Throughput was fine for a pure-Python data path: 1,701 MiB/s through the proxy against 2,036 MiB/s directly to the same asyncio upstream on one core. This guide builds the proxy, gets shutdown right in both directions, and adds the limits a proxy in production needs.

Prerequisites

1. Copy both directions concurrently

Each accepted connection gets an upstream connection and two copy tasks, one per direction. Each copy awaits drain() so a slow receiver slows the sender — backpressure end to end:

import asyncio

UPSTREAM = ("10.0.0.5", 5432)


async def pipe(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    try:
        while data := await reader.read(65536):
            writer.write(data)
            await writer.drain()                     # wait if the other side is slow
        if writer.can_write_eof():
            writer.write_eof()                       # forward the half-close (step 2)
    except ConnectionError:
        writer.close()


async def handle(client_reader, client_writer) -> None:
    try:
        upstream_reader, upstream_writer = await asyncio.open_connection(*UPSTREAM)
    except OSError:
        client_writer.close()                        # upstream down: refuse cleanly
        return
    try:
        async with asyncio.TaskGroup() as tg:
            tg.create_task(pipe(client_reader, upstream_writer))
            tg.create_task(pipe(upstream_reader, client_writer))
    finally:
        client_writer.close()
        upstream_writer.close()

The proxy holds at most one 64 KiB chunk per direction in Python, plus what the transports buffer up to their high-water marks — so memory per connection is bounded regardless of how much data flows. Measured, 1 GiB moved at 1,701 MiB/s through the proxy versus 2,036 MiB/s directly, with the proxy, upstream and sender sharing one machine.

Verify: curl --proxy or a database client through the proxy behaves identically to a direct connection.

A proxied connection with half-close A sequence of 8 messages between 3 participants. A proxied connection with half-close client proxy upstream connect open_connection request bytes pipe -> write + drain EOF (shutdown write) write_eof() reply bytes, then EOF pipe back, write_eof() Each direction ends independently; the connection ends when both have.

2. Forward half-close instead of closing

TCP connections are two independent byte streams. A client can finish sending (half-close) and still wait to receive — HTTP/1.0 clients, nc -N, many RPC and database protocols do this. A proxy that treats the first end-of-file as "connection over" cuts off the reply:

# Wrong: closing the upstream when the client stops sending
async def pipe_naive(reader, writer):
    while data := await reader.read(65536):
        writer.write(data)
        await writer.drain()
    writer.close()          # measured: the client got b"" instead of the upstream's reply

# Right: signal end-of-file and keep the other direction open
async def pipe(reader, writer):
    while data := await reader.read(65536):
        writer.write(data)
        await writer.drain()
    if writer.can_write_eof():
        writer.write_eof()  # measured: the client got b"got 100000\n"

Tested against an upstream that counts bytes until end-of-file and then replies: the naive proxy returned an empty response, the half-close proxy returned the reply. write_eof() sends a TCP FIN on that direction only; the TaskGroup keeps the connection open until both copies have finished, and the finally block then closes both sockets.

Verify: printf 'data' | nc -N proxyhost port against an upstream that replies after EOF prints the reply.

3. Handle errors in one direction

If either side resets the connection, the other copy task would otherwise wait forever on a read that will never complete. The TaskGroup handles that only if errors escape as exceptions — so let resets propagate, or cancel the sibling explicitly:

async def handle(client_reader, client_writer) -> None:
    upstream_reader, upstream_writer = await asyncio.open_connection(*UPSTREAM)
    to_upstream = asyncio.create_task(pipe(client_reader, upstream_writer))
    to_client = asyncio.create_task(pipe(upstream_reader, client_writer))
    try:
        done, pending = await asyncio.wait({to_upstream, to_client},
                                           return_when=asyncio.FIRST_EXCEPTION)
        if any(t.exception() for t in done if not t.cancelled()):
            for t in pending:
                t.cancel()                          # one side failed: stop the other
        await asyncio.gather(*pending, return_exceptions=True)
    finally:
        client_writer.transport.abort()             # reset both: do not flush to a broken peer
        upstream_writer.transport.abort()

On a clean finish both tasks complete normally; on a reset, the failing direction raises, the other is cancelled, and both sockets are aborted. Aborting rather than closing matters here: close() would try to flush buffered data to a peer that is already gone. With the TaskGroup version from step 1, the same happens automatically when pipe lets ConnectionError propagate instead of catching it.

Verify: kill the upstream mid-transfer; the client's connection is reset promptly and the proxy holds no tasks for it afterwards.

How each way of ending affects the other direction A grid of 4 rows by 3 columns. How each way of ending affects the other direction event this direction other direction clean EOF write_eof() to peer keeps flowing reset / error stops cancelled, both aborted idle timeout aborted aborted proxy shutdown finish, then abort after grace same Only a clean EOF leaves the other direction running.

4. Bound connections and idle time

A production proxy limits concurrent connections and reclaims idle ones, because each holds two file descriptors and two tasks:

MAX_CONNECTIONS = 5_000
IDLE_TIMEOUT = 300.0
slots = asyncio.Semaphore(MAX_CONNECTIONS)


async def pipe(reader, writer, activity: list[float]) -> None:
    loop = asyncio.get_running_loop()
    while True:
        async with asyncio.timeout(IDLE_TIMEOUT):
            data = await reader.read(65536)
        if not data:
            break
        activity[0] = loop.time()
        writer.write(data)
        await writer.drain()
    if writer.can_write_eof():
        writer.write_eof()


async def handle(client_reader, client_writer) -> None:
    if slots.locked():
        client_writer.transport.abort()               # at capacity
        return
    async with slots:
        ...                                           # as in step 1

The idle timeout here is per direction, which suits most protocols; for long-lived connections where one side is legitimately silent for long periods (a database client waiting for LISTEN notifications), time out only when both directions have been idle, using the shared activity timestamp. Also put a deadline on open_connection to the upstream, so a black-holed upstream does not hold client connections in limbo.

Verify: the proxy refuses connections beyond the limit and closes connections idle for longer than the timeout.

5. Use the proxy for testing

A small asyncio proxy is a powerful test tool: put it between a client and a real dependency and inject faults — latency, bandwidth limits, dropped connections, silent stalls — that are hard to produce otherwise:

class FaultyPipe:
    def __init__(self, delay: float = 0.0, drop_after: int | None = None, stall: bool = False):
        self.delay, self.drop_after, self.stall = delay, drop_after, stall

    async def __call__(self, reader, writer) -> None:
        sent = 0
        while data := await reader.read(65536):
            if self.stall:
                continue                               # middlebox that forgot the connection
            await asyncio.sleep(self.delay)            # added latency per chunk
            if self.drop_after is not None and sent + len(data) > self.drop_after:
                writer.transport.abort()               # mid-stream reset
                return
            writer.write(data)
            await writer.drain()
            sent += len(data)

The silent-stall mode is exactly how the stale-connection hang was reproduced in handling stale pooled connections after idle timeouts. Tools like Toxiproxy do the same at larger scale; an in-process version runs inside pytest with no extra services.

Verify: each fault mode produces the expected client-side error in a test, and the client's timeout and retry logic handles it.

Is a hand-written asyncio proxy the right tool? A decision on What is the proxy for with 3 outcomes. Is a hand-written asyncio proxy the right tool? What is the proxy for? tests: latency, resets, stalls asyncio proxy in-process, scriptable inspect or route by protocol asyncio + parser logic in Python production LB at high bandwidth HAProxy, Envoy, nginx kernel-level speed Python proxies shine where logic matters more than raw bandwidth.

Verification

A TCP proxy is correct when:

  • Both directions copy concurrently with drain() for backpressure.
  • Half-close is forwarded with write_eof(), and the connection ends when both directions have.
  • Errors in one direction stop the other and abort both sockets.
  • Connections, idle time and upstream connects are bounded.

Diagnostic Hook: count active connections, bytes per direction and how each connection ended (clean, reset, idle timeout, upstream connect failure). Connections that never end are a missing timeout; many empty responses to clients that half-close mean the proxy is closing on first EOF.

Pitfalls & edge cases

  • Closing on the first EOF. Tested: the client got an empty response.
  • No drain(). A fast sender and slow receiver buffer unbounded data in the proxy.
  • One stuck direction. Without cancellation, the copy task waits forever.
  • close() after a reset. It tries to flush to a dead peer; abort instead.

Frequently Asked Questions

How do I write a TCP proxy with asyncio?

For each accepted connection, open an upstream connection and run two tasks that read from one side and write to the other with await drain(), forwarding end-of-file with write_eof(), then close both connections when both tasks finish.

Why does my asyncio proxy cut off responses?

It probably closes both connections when one side finishes sending. Forward the half-close with writer.write_eof() instead; in testing, that turned an empty response into the full reply.

How fast is a TCP proxy written in Python asyncio?

In testing it moved 1,701 MiB/s on loopback against 2,036 MiB/s directly, enough for most uses other than high-bandwidth load balancing.

How do I simulate network failures in tests with asyncio?

Put a small asyncio proxy between client and server that adds delays, resets connections after a number of bytes, or silently stops forwarding.