Skip to content

Commit

Permalink
Add support in socks5 proxy
Browse files Browse the repository at this point in the history
  • Loading branch information
zaif-yuval committed Jan 1, 2025
1 parent 8d27699 commit 768bf17
Show file tree
Hide file tree
Showing 9 changed files with 69 additions and 3 deletions.
25 changes: 24 additions & 1 deletion pymongo/asynchronous/pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,11 @@ def _set_non_inheritable_non_atomic(fd: int) -> None:
def _set_non_inheritable_non_atomic(fd: int) -> None: # noqa: ARG001
"""Dummy function for platforms that don't provide fcntl."""

try:
from python_socks.sync import Proxy
from python_socks import ProxyType
except ImportError:
Proxy = ProxyType = None

_IS_SYNC = False

Expand Down Expand Up @@ -838,7 +843,25 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket
sock.settimeout(timeout)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, True)
_set_keepalive_times(sock)
sock.connect(sa)
if proxy := options.proxy:
if Proxy is None:
raise RuntimeError(
"In order to use SOCKS5 proxy, python_socks must be installed. "
"This can be done by re-installing pymongo with `pip install pymongo[socks]`"
)
proxy_host = proxy['host']
proxy_port = proxy['port'] or 1080
sock.connect((proxy_host, proxy_port))
proxy = Proxy(
ProxyType.SOCKS5,
proxy_host,
proxy_port,
proxy['username'],
proxy['password']
)
proxy.connect(sa[0], dest_port=sa[1], _socket=sock)
else:
sock.connect(sa)
return sock
except OSError as e:
err = e
Expand Down
1 change: 1 addition & 0 deletions pymongo/asynchronous/topology.py
Original file line number Diff line number Diff line change
Expand Up @@ -967,6 +967,7 @@ def _create_pool_for_monitor(self, address: _Address) -> Pool:
driver=options.driver,
pause_enabled=False,
server_api=options.server_api,
proxy=options.proxy,
)

return self._settings.pool_class(
Expand Down
10 changes: 10 additions & 0 deletions pymongo/client_options.py
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,15 @@ def _parse_pool_options(
ssl_context, tls_allow_invalid_hostnames = _parse_ssl_options(options)
load_balanced = options.get("loadbalanced")
max_connecting = options.get("maxconnecting", common.MAX_CONNECTING)
if proxy_host := options.get("proxyHost"):
proxy = {
"host": proxy_host,
"port": options.get("proxyPort"),
"username": options.get("proxyUserName"),
"password": options.get("proxyPassword"),
}
else:
proxy = None
return PoolOptions(
max_pool_size,
min_pool_size,
Expand All @@ -188,6 +197,7 @@ def _parse_pool_options(
load_balanced=load_balanced,
credentials=credentials,
is_sync=is_sync,
proxy=proxy,
)


Expand Down
4 changes: 4 additions & 0 deletions pymongo/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -729,6 +729,10 @@ def validate_server_monitoring_mode(option: str, value: str) -> str:
"srvmaxhosts": validate_non_negative_integer,
"timeoutms": validate_timeoutms,
"servermonitoringmode": validate_server_monitoring_mode,
"proxyhost": validate_string,
"proxyport": validate_positive_integer_or_none,
"proxyusername": validate_string_or_none,
"proxypassword": validate_string_or_none,
}

# Dictionary where keys are the names of URI options specific to pymongo,
Expand Down
9 changes: 8 additions & 1 deletion pymongo/pool_options.py
Original file line number Diff line number Diff line change
Expand Up @@ -312,6 +312,7 @@ class PoolOptions:
"__server_api",
"__load_balanced",
"__credentials",
"__proxy",
)

def __init__(
Expand All @@ -334,6 +335,7 @@ def __init__(
load_balanced: Optional[bool] = None,
credentials: Optional[MongoCredential] = None,
is_sync: Optional[bool] = True,
proxy: Optional[dict] = None,
):
self.__max_pool_size = max_pool_size
self.__min_pool_size = min_pool_size
Expand All @@ -353,7 +355,7 @@ def __init__(
self.__load_balanced = load_balanced
self.__credentials = credentials
self.__metadata = copy.deepcopy(_METADATA)

self.__proxy = copy.deepcopy(proxy)
if appname:
self.__metadata["application"] = {"name": appname}

Expand Down Expand Up @@ -522,3 +524,8 @@ def server_api(self) -> Optional[ServerApi]:
def load_balanced(self) -> Optional[bool]:
"""True if this Pool is configured in load balanced mode."""
return self.__load_balanced

@property
def proxy(self) -> Optional[dict]:
"""Proxy settings, if configured"""
return self.__proxy
20 changes: 19 additions & 1 deletion pymongo/synchronous/pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -836,7 +836,25 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket
sock.settimeout(timeout)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, True)
_set_keepalive_times(sock)
sock.connect(sa)
if proxy := options.proxy:
if Proxy is None:
raise RuntimeError(
"In order to use SOCKS5 proxy, python_socks must be installed. "
"This can be done by re-installing pymongo with `pip install pymongo[socks]`"
)
proxy_host = proxy['host']
proxy_port = proxy['port'] or 1080
sock.connect((proxy_host, proxy_port))
proxy = Proxy(
ProxyType.SOCKS5,
proxy_host,
proxy_port,
proxy['username'],
proxy['password']
)
proxy.connect(sa[0], dest_port=sa[1], _socket=sock)
else:
sock.connect(sa)
return sock
except OSError as e:
err = e
Expand Down
1 change: 1 addition & 0 deletions pymongo/synchronous/topology.py
Original file line number Diff line number Diff line change
Expand Up @@ -965,6 +965,7 @@ def _create_pool_for_monitor(self, address: _Address) -> Pool:
driver=options.driver,
pause_enabled=False,
server_api=options.server_api,
proxy=options.proxy,
)

return self._settings.pool_class(
Expand Down
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ ocsp = ["requirements/ocsp.txt"]
snappy = ["requirements/snappy.txt"]
test = ["requirements/test.txt"]
zstd = ["requirements/zstd.txt"]
socks = ["requirements/socks.txt"]

[tool.pytest.ini_options]
minversion = "7"
Expand Down
1 change: 1 addition & 0 deletions requirements/socks.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
python-socks[asyncio]

0 comments on commit 768bf17

Please sign in to comment.