From dbc067d0ff0711b2189b981f1556ebfe82c21974 Mon Sep 17 00:00:00 2001 From: RTUnreal Date: Tue, 25 Jun 2024 12:46:26 +0200 Subject: [PATCH] Fix unstable packet reading --- python_socks/_connectors/socks5_async.py | 12 +++++++++++- python_socks/_connectors/socks5_sync.py | 12 +++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/python_socks/_connectors/socks5_async.py b/python_socks/_connectors/socks5_async.py index ffab73f..b27dabf 100644 --- a/python_socks/_connectors/socks5_async.py +++ b/python_socks/_connectors/socks5_async.py @@ -3,6 +3,7 @@ from .._abc import AsyncSocketStream, AsyncResolver from .abc import AsyncConnector +from .._protocols.errors import ReplyError from .._protocols import socks5 from .._helpers import is_ip_address @@ -65,6 +66,15 @@ async def connect( data = conn.send(request) await stream.write_all(data) - data = await stream.read() + data = await stream.read_exact(4) + if data[3] == socks5.AddressType.IPV4: + data += await stream.read_exact(6) + elif data[3] == socks5.AddressType.IPV6: + data += await stream.read_exact(18) + elif data[3] == socks5.AddressType.DOMAIN: + data += await stream.read_exact(int.from_bytes(await stream.read_exact(1)) + 2) + else: # pragma: no cover + raise ReplyError(f'Invalid address type: {data[3]:#02X}') + reply: socks5.ConnectReply = conn.receive(data) return reply diff --git a/python_socks/_connectors/socks5_sync.py b/python_socks/_connectors/socks5_sync.py index 0d5b42d..05a9b04 100644 --- a/python_socks/_connectors/socks5_sync.py +++ b/python_socks/_connectors/socks5_sync.py @@ -3,6 +3,7 @@ from .._abc import SyncSocketStream, SyncResolver from .abc import SyncConnector +from .._protocols.errors import ReplyError from .._protocols import socks5 from .._helpers import is_ip_address @@ -56,6 +57,15 @@ def connect( data = conn.send(request) stream.write_all(data) - data = stream.read() + data = stream.read_exact(4) + if data[3] == socks5.AddressType.IPV4: + data += stream.read_exact(6) + elif data[3] == socks5.AddressType.IPV6: + data += stream.read_exact(18) + elif data[3] == socks5.AddressType.DOMAIN: + data += stream.read_exact(int.from_bytes(stream.read_exact(1)) + 2) + else: # pragma: no cover + raise ReplyError(f'Invalid address type: {data[3]:#02X}') + reply: socks5.ConnectReply = conn.receive(data) return reply