SNOW-694457: env-vars-proxy-leaking (#2451)

Co-authored-by: Berkay Öztürk <[email protected]>

(cherry picked from commit 87c623e)

Author: sfc-gh-fpawlowski

src/snowflake/connector/auth/_oauth_base.py

@@ -20,9 +20,12 @@
 )
 from ..errors import Error, ProgrammingError
 from ..network import OAUTH_AUTHENTICATOR
+from ..proxy import get_proxy_url
 from ..secret_detector import SecretDetector
 from ..token_cache import TokenCache, TokenKey, TokenType
 from ..vendored import urllib3
+from ..vendored.requests.utils import get_environ_proxies, select_proxy
+from ..vendored.urllib3.poolmanager import ProxyManager
 from .by_plugin import AuthByPlugin, AuthType
 
 if TYPE_CHECKING:
@@ -319,7 +322,13 @@ def _get_refresh_token_response(
             fields["scope"] = self._scope
         try:
             # TODO(SNOW-2229411) Session manager should be used here. It may require additional security validation (since we would transition from PoolManager to requests.Session) and some parameters would be passed implicitly. OAuth token exchange must NOT reuse pooled HTTP sessions. We should create a fresh SessionManager with use_pooling=False for each call.
-            return urllib3.PoolManager().request_encode_body(
+            proxy_url = self._resolve_proxy_url(conn, self._token_request_url)
+            http_client = (
+                ProxyManager(proxy_url=proxy_url)
+                if proxy_url
+                else urllib3.PoolManager()
+            )
+            return http_client.request_encode_body(
                 "POST",
                 self._token_request_url,
                 encode_multipart=False,
@@ -359,7 +368,11 @@ def _get_request_token_response(
         fields: dict[str, str],
     ) -> (str | None, str | None):
         # TODO(SNOW-2229411) Session manager should be used here. It may require additional security validation (since we would transition from PoolManager to requests.Session) and some parameters would be passed implicitly. Token request must bypass HTTP connection pools.
-        resp = urllib3.PoolManager().request_encode_body(
+        proxy_url = self._resolve_proxy_url(connection, self._token_request_url)
+        http_client = (
+            ProxyManager(proxy_url=proxy_url) if proxy_url else urllib3.PoolManager()
+        )
+        resp = http_client.request_encode_body(
             "POST",
             self._token_request_url,
             headers=self._create_token_request_headers(),
@@ -400,3 +413,25 @@ def _create_token_request_headers(self) -> dict[str, str]:
             "Accept": "application/json",
             "Content-Type": "application/x-www-form-urlencoded; charset=UTF-8",
         }
+
+    @staticmethod
+    def _resolve_proxy_url(
+        connection: SnowflakeConnection, request_url: str
+    ) -> str | None:
+        # TODO(SNOW-2229411) Session manager should be used instead. It may require additional security validation.
+        """Resolve proxy URL from explicit config first, then environment variables."""
+        # First try explicit proxy configuration from connection parameters
+        proxy_url = get_proxy_url(
+            connection.proxy_host,
+            connection.proxy_port,
+            connection.proxy_user,
+            connection.proxy_password,
+        )
+
+        if proxy_url:
+            return proxy_url
+
+        # Fall back to environment variables (HTTP_PROXY, HTTPS_PROXY)
+        # Use proper proxy selection that considers the URL scheme
+        proxies = get_environ_proxies(request_url)
+        return select_proxy(request_url, proxies)

src/snowflake/connector/auth/oauth_code.py

@@ -269,6 +269,7 @@ def _do_authorization_request(
             "login. If you can't see it, check existing browser windows, "
             "or your OS settings. Press CTRL+C to abort and try again..."
         )
+        # TODO(SNOW-2229411) Investigate if Session manager / Http Config should be used here.
         code, state = (
             self._receive_authorization_callback(callback_server, connection)
             if webbrowser.open(authorization_request)

src/snowflake/connector/connection.py

@@ -27,7 +27,7 @@
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
 
-from . import errors, proxy
+from . import errors
 from ._query_context_cache import QueryContextCache
 from ._utils import (
     _DEFAULT_VALUE_SERVER_DOP_CAP_FOR_FILE_TRANSFER,
@@ -924,6 +924,10 @@ def connect(self, **kwargs) -> None:
         self._http_config = HttpConfig(
             adapter_factory=ProxySupportAdapterFactory(),
             use_pooling=(not self.disable_request_pooling),
+            proxy_host=self.proxy_host,
+            proxy_port=self.proxy_port,
+            proxy_user=self.proxy_user,
+            proxy_password=self.proxy_password,
         )
         self._session_manager = SessionManager(self._http_config)
 
@@ -1125,10 +1129,6 @@ 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,

src/snowflake/connector/proxy.py

@@ -1,43 +1,28 @@
 #!/usr/bin/env python
 from __future__ import annotations
 
-import os
 
-
-def set_proxies(
+def get_proxy_url(
     proxy_host: str | None,
     proxy_port: str | None,
     proxy_user: str | None = None,
     proxy_password: str | None = None,
-) -> dict[str, str] | None:
-    """Sets proxy dict for requests."""
-    PREFIX_HTTP = "http://"
-    PREFIX_HTTPS = "https://"
-    proxies = None
+) -> str | None:
+    http_prefix = "http://"
+    https_prefix = "https://"
+
     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}"
+
+    return None

src/snowflake/connector/result_batch.py

@@ -26,7 +26,7 @@
 from .options import installed_pandas
 from .options import pyarrow as pa
 from .secret_detector import SecretDetector
-from .session_manager import SessionManager
+from .session_manager import HttpConfig, SessionManager
 from .time_util import TimerContextManager
 
 logger = getLogger(__name__)
@@ -261,6 +261,8 @@ def __init__(
             [s._to_result_metadata_v1() for s in schema] if schema is not None else None
         )
         self._use_dict_result = use_dict_result
+        # Passed to contain the configured Http behavior in case the connectio is no longer active for the download
+        # Can be overridden with setters if needed.
         self._session_manager = session_manager
         self._metrics: dict[str, int] = {}
         self._data: str | list[tuple[Any, ...]] | None = None
@@ -300,6 +302,25 @@ def uncompressed_size(self) -> int | None:
     def column_names(self) -> list[str]:
         return [col.name for col in self._schema]
 
+    @property
+    def session_manager(self) -> SessionManager | None:
+        return self._session_manager
+
+    @session_manager.setter
+    def session_manager(self, session_manager: SessionManager | None) -> None:
+        self._session_manager = session_manager
+
+    @property
+    def http_config(self):
+        return self._session_manager.config
+
+    @http_config.setter
+    def http_config(self, config: HttpConfig) -> None:
+        if self._session_manager:
+            self._session_manager.config = config
+        else:
+            self._session_manager = SessionManager(config=config)
+
     def __iter__(
         self,
     ) -> Iterator[dict | Exception] | Iterator[tuple | Exception]:

src/snowflake/connector/session_manager.py

@@ -10,6 +10,7 @@
 from typing import TYPE_CHECKING, Any, Callable, Generator, Generic, Mapping, TypeVar
 
 from .compat import urlparse
+from .proxy import get_proxy_url
 from .vendored import requests
 from .vendored.requests import Response, Session
 from .vendored.requests.adapters import BaseAdapter, HTTPAdapter
@@ -79,8 +80,15 @@ def get_connection(
             proxy_manager = self.proxy_manager_for(proxy)
 
             if isinstance(proxy_manager, ProxyManager):
-                # Add Host to proxy header SNOW-232777
-                proxy_manager.proxy_headers["Host"] = parsed_url.hostname
+                # Add Host to proxy header SNOW-232777 and SNOW-694457
+
+                # RFC 7230 / 5.4 – a proxy’s Host header must repeat the request authority
+                # verbatim: <hostname>[:<port>] with IPv6 still in [brackets].  We take that
+                # straight from urlparse(url).netloc, which preserves port and brackets (and case-sensitive hostname).
+                # Note: netloc also keeps user-info (user:pass@host) if present in URL. The driver never sends
+                # URLs with embedded credentials, so we leave them unhandled — for full support
+                # we’d need to manually concatenate hostname with optional port and IPv6 brackets.
+                proxy_manager.proxy_headers["Host"] = parsed_url.netloc
             else:
                 logger.debug(
                     f"Unable to set 'Host' to proxy manager of type {type(proxy_manager)} as"
@@ -112,6 +120,10 @@ class BaseHttpConfig:
 
     use_pooling: bool = True
     max_retries: int | None = REQUESTS_RETRY
+    proxy_host: str | None = None
+    proxy_port: str | None = None
+    proxy_user: str | None = None
+    proxy_password: str | None = None
 
     def copy_with(self, **overrides: Any) -> BaseHttpConfig:
         """Return a new config with overrides applied."""
@@ -325,13 +337,13 @@ class SessionManager(_RequestVerbsUsingSessionMixin):
     **Two Operating Modes**:
     - use_pooling=False: One-shot sessions (create, use, close) - suitable for infrequent requests
     - use_pooling=True: Per-hostname session pools - reuses TCP connections, avoiding handshake
-      and SSL/TLS negotiation overhead for repeated requests to the same host
+      and SSL/TLS negotiation overhead for repeated requests to the same host.
 
     **Key Benefits**:
     - Centralized HTTP configuration management and easy propagation across the codebase
     - Consistent proxy setup (SNOW-694457) and headers customization (SNOW-2043816)
     - HTTPAdapter customization for connection-level request manipulation
-    - Performance optimization through connection reuse for high-traffic scenarios
+    - Performance optimization through connection reuse for high-traffic scenarios.
 
     **Usage**: Create the base session manager, then use clone() for derived managers to ensure
     proper config propagation. Pre-commit checks enforce usage to prevent code drift back to
@@ -347,7 +359,6 @@ def __init__(self, config: HttpConfig | None = None, **http_config_kwargs) -> No
             logger.debug("Creating a config for the SessionManager")
             config = HttpConfig(**http_config_kwargs)
         self._cfg: HttpConfig = config
-
         self._sessions_map: dict[str | None, SessionPool] = collections.defaultdict(
             lambda: SessionPool(self)
         )
@@ -370,6 +381,19 @@ def from_config(cls, cfg: HttpConfig, **overrides: Any) -> SessionManager:
     def config(self) -> HttpConfig:
         return self._cfg
 
+    @config.setter
+    def config(self, cfg: HttpConfig) -> None:
+        self._cfg = cfg
+
+    @property
+    def proxy_url(self) -> str:
+        return get_proxy_url(
+            self._cfg.proxy_host,
+            self._cfg.proxy_port,
+            self._cfg.proxy_user,
+            self._cfg.proxy_password,
+        )
+
     @property
     def use_pooling(self) -> bool:
         return self._cfg.use_pooling
@@ -427,6 +451,7 @@ def _mount_adapters(self, session: requests.Session) -> None:
     def make_session(self) -> Session:
         session = requests.Session()
         self._mount_adapters(session)
+        session.proxies = {"http": self.proxy_url, "https": self.proxy_url}
         return session
 
     @contextlib.contextmanager

test/conftest.py

@@ -5,6 +5,8 @@
 from contextlib import contextmanager
 from logging import getLogger
 from pathlib import Path
+from test.test_utils.cross_module_fixtures.http_fixtures import *  # NOQA
+from test.test_utils.cross_module_fixtures.wiremock_fixtures import *  # NOQA
 from typing import Generator
 
 import pytest

test/data/wiremock/mappings/auth/password/successful_flow.json

@@ -0,0 +1,60 @@
+{
+  "mappings": [
+    {
+      "request": {
+        "urlPathPattern": "/session/v1/login-request.*",
+        "method": "POST",
+        "bodyPatterns": [
+          {
+            "equalToJson" : {
+              "data": {
+                "LOGIN_NAME": "testUser",
+                "PASSWORD": "testPassword"
+              }
+            },
+            "ignoreExtraElements" : true
+          }
+        ]
+      },
+      "response": {
+        "status": 200,
+        "jsonBody": {
+          "data": {
+            "masterToken": "master token",
+            "token": "session token",
+            "validityInSeconds": 3600,
+            "masterValidityInSeconds": 14400,
+            "displayUserName": "TEST_USER",
+            "serverVersion": "8.48.0 b2024121104444034239f05",
+            "firstLogin": false,
+            "remMeToken": null,
+            "remMeValidityInSeconds": 0,
+            "healthCheckInterval": 45,
+            "newClientForUpgrade": "3.12.3",
+            "sessionId": 1172562260498,
+            "parameters": [
+              {
+                "name": "CLIENT_PREFETCH_THREADS",
+                "value": 4
+              }
+            ],
+            "sessionInfo": {
+              "databaseName": "TEST_DB",
+              "schemaName": "TEST_GO",
+              "warehouseName": "TEST_XSMALL",
+              "roleName": "ANALYST"
+            },
+            "idToken": null,
+            "idTokenValidityInSeconds": 0,
+            "responseData": null,
+            "mfaToken": null,
+            "mfaTokenValidityInSeconds": 0
+          },
+          "code": null,
+          "message": null,
+          "success": true
+        }
+      }
+    }
+  ]
+}

test/data/wiremock/mappings/generic/proxy_forward_all.json

@@ -0,0 +1,12 @@
+{
+  "request": {
+    "urlPattern": "/.*",
+    "method": "ANY"
+  },
+  "response": {
+    "proxyBaseUrl": "{{TARGET_HTTP_HOST_WITH_PORT}}",
+    "additionalProxyRequestHeaders": {
+        "Via": "1.1 wiremock-proxy"
+    }
+  }
+}