# Love you njsmith <3
import socket
import structio


async def proxy_one_way(source, sink: structio.AsyncSocket):
    while True:
        data = await source.receive(1024)
        if not data:
            await sink.shutdown(socket.SHUT_WR)
            break
        await sink.send_all(data)


async def proxy_two_way(a, b: structio.AsyncSocket):
    async with structio.create_pool() as pool:
        pool.spawn(proxy_one_way, a, b)
        pool.spawn(proxy_one_way, b, a)


async def main():
    with structio.skip_after(10):
        a = await structio.socket.connect_tcp_socket("localhost", 12345)
        b = await structio.socket.connect_tcp_socket("localhost", 54321)
        async with a, b:
            await proxy_two_way(a, b)
    print("all done!")


structio.run(main)