romis2012 / python-socks

Core proxy client (SOCKS4, SOCKS5, HTTP) functionality for Python
Apache License 2.0
101 stars 18 forks source link

`ConnectionResetError('Connection lost')` after .drain()'ing #29

Closed Newcool1230 closed 9 months ago

Newcool1230 commented 9 months ago

Hey there,

I'm having some trouble with async proxy connections. My goal is to create an asyncio HTTP proxy server. I'm getting a connection lost error after trying to forward the data.

Task exception was never retrieved
future: <Task finished name='Task-5' coro=<handle_client() done, defined at \v2\testing.py:14> exception=ConnectionResetError('Connection lost')>
Traceback (most recent call last):
  File "\v2\testing.py", line 40, in handle_client
    await asyncio.gather(
  File "\v2\testing.py", line 12, in forward_data
    await dest.drain()
  File "\Python311\Lib\asyncio\streams.py", line 377, in drain
    await self._protocol._drain_helper()
  File "\Python311\Lib\asyncio\streams.py", line 166, in _drain_helper
    raise ConnectionResetError('Connection lost')
ConnectionResetError: Connection lost

I'm testing this through Firefox's proxy.onRequest API:

// background.js
browser.proxy.onRequest.addListener((requestInfo) => {
    console.log(requestInfo);
    if (requestInfo.url.includes("example.com") || requestInfo.url.includes("ifconfig.me")) {
        console.log(`Proxying: ${requestInfo.url}`);
        return { type: "http", host: "127.0.0.1", port: 2222 };
    }
    return { type: "direct" };
}, { urls: ["<all_urls>"], });

manifest.json

{
    "manifest_version": 2,
    "name": "TTP Proxy Test",
    "author": "ME",
    "homepage_url": "https://127.0.0.1",
    "version": "2024.1.19",
    "description": "testing extension",
    "browser_specific_settings": {
        "gecko": {
          "id": "testing@me.com"
        }
    },
    "background": {
        "scripts": ["background.js"]
    },
    "permissions": [
        "proxy",
        "<all_urls>"
    ]
}
# proxy.py
import asyncio
import ssl
from asyncio.streams import StreamReader, StreamWriter
from python_socks.async_.asyncio import Proxy

async def forward_data(src: StreamReader, dest: StreamWriter):
    while True:
        data = await src.read(4096)
        if not data:
            break
        dest.write(data)
        await dest.drain()

async def handle_client(client_reader: StreamReader, client_writer: StreamWriter):
    request = await client_reader.read(4096)
    print(request)
    first_line = request.decode().split('\n')[0]
    method, url, protocol = first_line.split()
    print(first_line)

    if method == 'CONNECT':
        target_host, target_port = url.split(':')
        target_port = int(target_port)
        print(target_host, target_port)

        proxy = Proxy.from_url('socks5://user:password@127.0.0.1:1080')
        sock = await proxy.connect(dest_host=target_host, dest_port=target_port)

        target_reader, target_writer = await asyncio.open_connection(
            host=None,
            port=None,
            sock=sock,
            ssl=ssl.create_default_context(),
            server_hostname=target_host,
        )

        client_writer.write(b'HTTP/1.1 200 Connection established\r\n\r\n')
        await client_writer.drain()

        await asyncio.gather(
            forward_data(client_reader, target_writer),
            forward_data(target_reader, client_writer)
        )
    else:
        print('Ignoring non CONNECT requests')

async def main():
    server = await asyncio.start_server(
        handle_client, '127.0.0.1', 2222
    )
    print(f'Serving on {server.sockets[0].getsockname()}')

    async with server:
        await server.serve_forever()

if __name__ == "__main__":
    asyncio.run(main())

If I run it normally without the proxy connection it will work just fine.

 target_reader, target_writer = await asyncio.open_connection(target_host, target_port)

I was also able to get the proxy working in sync but once I moved to asyncio it will lose connection.

romis2012 commented 9 months ago

Connections are sometimes closed by the client, this is normal.

The forward_data function should look something like this

async def forward_data(reader: StreamReader, writer: StreamWriter):
    try:
        while not reader.at_eof() and not writer.is_closing():
            data = await reader.read(65536)
            if not data:
                break
            writer.write(data)
            await writer.drain()
    except (asyncio.CancelledError, ConnectionResetError):  # maybe some other exceptions
        pass  # log it?
    finally:
        writer.close()

Also you should not pass ssl and server_hostname arguments to asyncio.open_connection, TLS communication must be performed directly between the client and the target server

Newcool1230 commented 9 months ago

Also you should not pass ssl and server_hostname arguments to asyncio.open_connection, TLS communication must be performed directly between the client and the target server

Thank you for getting back to me so quickly. This was the answer! After removing ssl and server_hostname it works correctly!

target_reader, target_writer = await asyncio.open_connection(sock=sock)

I've also updated my forward_data to better catch these errors! Thank you!