From b94add99ea13fbf2fdf5d1b185f47f628055b203 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Berkay=20=C3=96zt=C3=BCrk?= Date: Sun, 15 Jan 2023 12:15:26 +0300 Subject: [PATCH] Fix leaking proxy env vars --- src/snowflake/connector/connection.py | 10 +++--- src/snowflake/connector/network.py | 7 ++++ src/snowflake/connector/proxy.py | 50 ++++++++++----------------- test/integ/test_connection.py | 10 +++--- test/unit/test_proxies.py | 32 ++++++----------- 5 files changed, 47 insertions(+), 62 deletions(-) diff --git a/src/snowflake/connector/connection.py b/src/snowflake/connector/connection.py index 7fc3d2a6b..95a584c74 100644 --- a/src/snowflake/connector/connection.py +++ b/src/snowflake/connector/connection.py @@ -22,7 +22,7 @@ from time import strptime from typing import Any, Callable, Generator, Iterable, NamedTuple, Sequence -from . import errors, proxy +from . import errors from .auth import ( FIRST_PARTY_AUTHENTICATORS, Auth, @@ -715,16 +715,16 @@ def __open_connection(self): use_numpy=self._numpy, support_negative_year=self._support_negative_year ) - proxy.set_proxies( - self.proxy_host, self.proxy_port, self.proxy_user, self.proxy_password - ) - self._rest = SnowflakeRestful( host=self.host, port=self.port, protocol=self._protocol, inject_client_pause=self._inject_client_pause, connection=self, + proxy_host=self.proxy_host, + proxy_port=self.proxy_port, + proxy_user=self.proxy_user, + proxy_password=self.proxy_password, ) logger.debug("REST API object was created: %s:%s", self.host, self.port) diff --git a/src/snowflake/connector/network.py b/src/snowflake/connector/network.py index e61e7d856..ee45d7934 100644 --- a/src/snowflake/connector/network.py +++ b/src/snowflake/connector/network.py @@ -75,6 +75,7 @@ ProgrammingError, ServiceUnavailableError, ) +from .proxy import get_proxy_url from .sqlstate import ( SQLSTATE_CONNECTION_NOT_EXISTS, SQLSTATE_CONNECTION_REJECTED, @@ -346,9 +347,14 @@ def __init__( protocol="http", inject_client_pause=0, connection: SnowflakeConnection | None = None, + proxy_host: str | None = None, + proxy_port: str | None = None, + proxy_user: str | None = None, + proxy_password: str | None = None, ): self._host = host self._port = port + self._proxy = get_proxy_url(proxy_host, proxy_port, proxy_user, proxy_password) self._protocol = protocol self._inject_client_pause = inject_client_pause self._connection = connection @@ -1130,6 +1136,7 @@ def make_requests_session(self): s.mount("http://", ProxySupportAdapter(max_retries=REQUESTS_RETRY)) s.mount("https://", ProxySupportAdapter(max_retries=REQUESTS_RETRY)) s._reuse_count = itertools.count() + s.proxies = {"http": self._proxy, "https": self._proxy} return s @contextlib.contextmanager diff --git a/src/snowflake/connector/proxy.py b/src/snowflake/connector/proxy.py index 884c0558f..08f6bd7fe 100644 --- a/src/snowflake/connector/proxy.py +++ b/src/snowflake/connector/proxy.py @@ -5,38 +5,26 @@ from __future__ import annotations -import os +def get_proxy_url( + proxy_host: str | None, + proxy_port: str | None, + proxy_user: str | None = None, + proxy_password: str | None = None, +) -> str | None: + http_prefix = "http://" + https_prefix = "https://" -def set_proxies(proxy_host, proxy_port, proxy_user=None, proxy_password=None): - """Sets proxy dict for requests.""" - PREFIX_HTTP = "http://" - PREFIX_HTTPS = "https://" - proxies = None if proxy_host and proxy_port: - if proxy_host.startswith(PREFIX_HTTP): - proxy_host = proxy_host[len(PREFIX_HTTP) :] - elif proxy_host.startswith(PREFIX_HTTPS): - proxy_host = proxy_host[len(PREFIX_HTTPS) :] - if proxy_user or proxy_password: - proxy_auth = "{proxy_user}:{proxy_password}@".format( - proxy_user=proxy_user if proxy_user is not None else "", - proxy_password=proxy_password if proxy_password is not None else "", - ) + if proxy_host.startswith(http_prefix): + host = proxy_host[len(http_prefix) :] + elif proxy_host.startswith(https_prefix): + host = proxy_host[len(https_prefix) :] else: - proxy_auth = "" - proxies = { - "http": "http://{proxy_auth}{proxy_host}:{proxy_port}".format( - proxy_host=proxy_host, - proxy_port=str(proxy_port), - proxy_auth=proxy_auth, - ), - "https": "http://{proxy_auth}{proxy_host}:{proxy_port}".format( - proxy_host=proxy_host, - proxy_port=str(proxy_port), - proxy_auth=proxy_auth, - ), - } - os.environ["HTTP_PROXY"] = proxies["http"] - os.environ["HTTPS_PROXY"] = proxies["https"] - return proxies + host = proxy_host + auth = ( + f"{proxy_user or ''}:{proxy_password or ''}@" + if proxy_user or proxy_password + else "" + ) + return f"{http_prefix}{auth}{host}:{proxy_port}" diff --git a/test/integ/test_connection.py b/test/integ/test_connection.py index 63fedd786..9c4f12145 100644 --- a/test/integ/test_connection.py +++ b/test/integ/test_connection.py @@ -433,6 +433,8 @@ def test_invalid_account_timeout(): @pytest.mark.timeout(15) def test_invalid_proxy(db_parameters): + http_proxy = os.environ.get("HTTP_PROXY") + https_proxy = os.environ.get("HTTPS_PROXY") with pytest.raises(OperationalError): snowflake.connector.connect( protocol="http", @@ -441,13 +443,13 @@ def test_invalid_proxy(db_parameters): password=db_parameters["password"], host=db_parameters["host"], port=db_parameters["port"], - login_timeout=5, + login_timeout=0, proxy_host="localhost", proxy_port="3333", ) - # NOTE environment variable is set if the proxy parameter is specified. - del os.environ["HTTP_PROXY"] - del os.environ["HTTPS_PROXY"] + # Proxy environment variables should not change + assert os.environ.get("HTTP_PROXY") == http_proxy + assert os.environ.get("HTTPS_PROXY") == https_proxy @pytest.mark.timeout(15) diff --git a/test/unit/test_proxies.py b/test/unit/test_proxies.py index 105ff4691..865c120b9 100644 --- a/test/unit/test_proxies.py +++ b/test/unit/test_proxies.py @@ -4,29 +4,17 @@ # from __future__ import annotations -import os +def test_get_proxy_url(): + from snowflake.connector.proxy import get_proxy_url -def test_set_proxies(): - from snowflake.connector.proxy import set_proxies + assert get_proxy_url("host", "port", "user", "password") == ( + "http://user:password@host:port" + ) + assert get_proxy_url("host", "port") == "http://host:port" - assert set_proxies("proxyhost", "8080") == { - "http": "http://proxyhost:8080", - "https": "http://proxyhost:8080", - } - assert set_proxies("http://proxyhost", "8080") == { - "http": "http://proxyhost:8080", - "https": "http://proxyhost:8080", - } - assert set_proxies("http://proxyhost", "8080", "testuser", "testpass") == { - "http": "http://testuser:testpass@proxyhost:8080", - "https": "http://testuser:testpass@proxyhost:8080", - } - assert set_proxies("proxyhost", "8080", "testuser", "testpass") == { - "http": "http://testuser:testpass@proxyhost:8080", - "https": "http://testuser:testpass@proxyhost:8080", - } + assert get_proxy_url("http://host", "port") == "http://host:port" - # NOTE environment variable is set if the proxy parameter is specified. - del os.environ["HTTP_PROXY"] - del os.environ["HTTPS_PROXY"] + assert get_proxy_url("https://host", "port", "user", "password") == ( + "http://user:password@host:port" + )