Skip to content

Commit

Permalink
Fix leaking proxy env vars
Browse files Browse the repository at this point in the history
  • Loading branch information
ozturkberkay committed Jan 15, 2023
1 parent e0d8456 commit b94add9
Show file tree
Hide file tree
Showing 5 changed files with 47 additions and 62 deletions.
10 changes: 5 additions & 5 deletions src/snowflake/connector/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)

Expand Down
7 changes: 7 additions & 0 deletions src/snowflake/connector/network.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@
ProgrammingError,
ServiceUnavailableError,
)
from .proxy import get_proxy_url
from .sqlstate import (
SQLSTATE_CONNECTION_NOT_EXISTS,
SQLSTATE_CONNECTION_REJECTED,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
50 changes: 19 additions & 31 deletions src/snowflake/connector/proxy.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}"
10 changes: 6 additions & 4 deletions test/integ/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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)
Expand Down
32 changes: 10 additions & 22 deletions test/unit/test_proxies.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

0 comments on commit b94add9

Please sign in to comment.