SNOW-2194055: Separate server and redirect URIs in AuthHttpServer

DESCRIPTION.md

@@ -14,6 +14,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Pin lower versions of dependencies to oldest version without vulnerabilities.
   - Added no_proxy parameter for proxy configuration without using environmental variables.
   - Added OAUTH_AUTHORIZATION_CODE and OAUTH_CLIENT_CREDENTIALS to list of authenticators that don't require user to be set
+  - Added `oauth_socket_uri` connection parameter allowing to separate server and redirect URIs for local OAuth server.
 
 - v4.0.0(October 09,2025)
   - Added support for checking certificates revocation using revocation lists (CRLs)

src/snowflake/connector/auth/_http_server.py

@@ -70,8 +70,10 @@ def __init__(
         self,
         uri: str,
         buf_size: int = 16384,
+        redirect_uri: str | None = None,
     ) -> None:
         parsed_uri = urllib.parse.urlparse(uri)
+        parsed_redirect = urllib.parse.urlparse(redirect_uri) if redirect_uri else None
         self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self.buf_size = buf_size
         if os.getenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", "False").lower() == "true":
@@ -82,30 +84,34 @@ def __init__(
             else:
                 self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
 
-        port = parsed_uri.port or 0
+        if parsed_redirect and self._is_local_uri(parsed_redirect):
+            server_port = parsed_redirect.port or 0
+        else:
+            server_port = parsed_uri.port or 0
+
         for attempt in range(1, self.DEFAULT_MAX_ATTEMPTS + 1):
             try:
                 self._socket.bind(
                     (
                         parsed_uri.hostname,
-                        port,
+                        server_port,
                     )
                 )
                 break
             except socket.gaierror as ex:
                 logger.error(
-                    f"Failed to bind authorization callback server to port {port}: {ex}"
+                    f"Failed to bind authorization callback server to port {server_port}: {ex}"
                 )
                 raise
             except OSError as ex:
                 if attempt == self.DEFAULT_MAX_ATTEMPTS:
                     logger.error(
-                        f"Failed to bind authorization callback server to port {port}: {ex}"
+                        f"Failed to bind authorization callback server to port {server_port}: {ex}"
                     )
                     raise
                 logger.warning(
                     f"Attempt {attempt}/{self.DEFAULT_MAX_ATTEMPTS}. "
-                    f"Failed to bind authorization callback server to port {port}: {ex}"
+                    f"Failed to bind authorization callback server to port {server_port}: {ex}"
                 )
                 time.sleep(self.PORT_BIND_TIMEOUT / self.PORT_BIND_MAX_ATTEMPTS)
         try:
@@ -114,16 +120,47 @@ def __init__(
             logger.error(f"Failed to start listening for auth callback: {ex}")
             self.close()
             raise
-        port = self._socket.getsockname()[1]
+
+        server_port = self._socket.getsockname()[1]
         self._uri = urllib.parse.ParseResult(
             scheme=parsed_uri.scheme,
-            netloc=parsed_uri.hostname + ":" + str(port),
+            netloc=parsed_uri.hostname + ":" + str(server_port),
             path=parsed_uri.path,
             params=parsed_uri.params,
             query=parsed_uri.query,
             fragment=parsed_uri.fragment,
         )
 
+        if parsed_redirect:
+            if (
+                self._is_local_uri(parsed_redirect)
+                and server_port != parsed_redirect.port
+            ):
+                logger.debug(
+                    f"Updating redirect port {parsed_redirect.port} to match the server port {server_port}."
+                )
+                self._redirect_uri = urllib.parse.ParseResult(
+                    scheme=parsed_redirect.scheme,
+                    netloc=parsed_redirect.hostname + ":" + str(server_port),
+                    path=parsed_redirect.path,
+                    params=parsed_redirect.params,
+                    query=parsed_redirect.query,
+                    fragment=parsed_redirect.fragment,
+                )
+            else:
+                self._redirect_uri = parsed_redirect
+        else:
+            # For backwards compatibility
+            self._redirect_uri = self._uri
+
+    @staticmethod
+    def _is_local_uri(uri):
+        return uri.hostname in ("localhost", "127.0.0.1")
+
+    @property
+    def redirect_uri(self) -> str | None:
+        return self._redirect_uri.geturl()
+
     @property
     def url(self) -> str:
         return self._uri.geturl()

src/snowflake/connector/auth/oauth_code.py

@@ -65,6 +65,7 @@ def __init__(
         external_browser_timeout: int | None = None,
         enable_single_use_refresh_tokens: bool = False,
         connection: SnowflakeConnection | None = None,
+        uri: str | None = None,
         **kwargs,
     ) -> None:
         authentication_url, redirect_uri = self._validate_oauth_code_uris(
@@ -92,6 +93,7 @@ def __init__(
         self._origin: str | None = None
         self._authentication_url = authentication_url
         self._redirect_uri = redirect_uri
+        self._uri = uri
         self._state = secrets.token_urlsafe(43)
         logger.debug("chose oauth state: %s", "".join("*" for _ in self._state))
         self._protocol = "http"
@@ -117,7 +119,10 @@ def _request_tokens(
     ) -> (str | None, str | None):
         """Web Browser based Authentication."""
         logger.debug("authenticating with OAuth authorization code flow")
-        with AuthHttpServer(self._redirect_uri) as callback_server:
+        with AuthHttpServer(
+            redirect_uri=self._redirect_uri,
+            uri=self._uri or self._redirect_uri,  # for backward compatibility
+        ) as callback_server:
             code = self._do_authorization_request(callback_server, conn)
             return self._do_token_request(code, callback_server, conn)
 
@@ -260,7 +265,7 @@ def _do_authorization_request(
         connection: SnowflakeConnection,
     ) -> str | None:
         authorization_request = self._construct_authorization_request(
-            callback_server.url
+            callback_server.redirect_uri
         )
         logger.debug("step 1: going to open authorization URL")
         print(
@@ -315,7 +320,7 @@ def _do_token_request(
         fields = {
             "grant_type": "authorization_code",
             "code": code,
-            "redirect_uri": callback_server.url,
+            "redirect_uri": callback_server.redirect_uri,
         }
         if self._enable_single_use_refresh_tokens:
             fields["enable_single_use_refresh_tokens"] = "true"

src/snowflake/connector/connection.py

@@ -384,6 +384,11 @@ def _get_private_bytes_from_file(
         # SNOW-1825621: OAUTH implementation
     ),
     "oauth_redirect_uri": ("http://127.0.0.1", str),
+    "oauth_socket_uri": (
+        "http://127.0.0.1",
+        str,
+        # SNOW-2194055: Separate server and redirect URIs in AuthHttpServer
+    ),
     "oauth_scope": (
         "",
         str,
@@ -1456,6 +1461,7 @@ def __open_connection(self):
                         host=self.host, port=self.port
                     ),
                     redirect_uri=self._oauth_redirect_uri,
+                    uri=self._oauth_socket_uri,
                     scope=self._oauth_scope,
                     pkce_enabled=not self._oauth_disable_pkce,
                     token_cache=(

src/snowflake/connector/constants.py

@@ -442,7 +442,6 @@ class IterUnit(Enum):
 ENV_VAR_PARTNER = "SF_PARTNER"
 ENV_VAR_TEST_MODE = "SNOWFLAKE_TEST_MODE"
 
-
 _DOMAIN_NAME_MAP = {_DEFAULT_HOSTNAME_TLD: "GLOBAL", _CHINA_HOSTNAME_TLD: "CHINA"}
 
 _CONNECTIVITY_ERR_MSG = (

test/unit/test_auth_callback_server.py

@@ -24,7 +24,9 @@ def test_auth_callback_success(monkeypatch, dontwait, timeout, reuse_port) -> No
     monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", reuse_port)
     monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT", dontwait)
     test_response: requests.Response | None = None
-    with AuthHttpServer("http://127.0.0.1/test_request") as callback_server:
+    with AuthHttpServer(
+        "http://127.0.0.1/test_request",
+    ) as callback_server:
 
         def request_callback():
             nonlocal test_response
@@ -57,7 +59,155 @@ def request_callback():
 def test_auth_callback_timeout(monkeypatch, dontwait, timeout, reuse_port) -> None:
     monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", reuse_port)
     monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT", dontwait)
-    with AuthHttpServer("http://127.0.0.1/test_request") as callback_server:
+    with AuthHttpServer(
+        "http://127.0.0.1/test_request",
+    ) as callback_server:
         block, client_socket = callback_server.receive_block(timeout=timeout)
         assert block is None
         assert client_socket is None
+
+
+@pytest.mark.parametrize(
+    "socket_host",
+    [
+        "127.0.0.1",
+        "localhost",
+    ],
+)
+@pytest.mark.parametrize(
+    "socket_port",
+    [
+        "",
+        ":0",
+        ":12345",
+    ],
+)
+@pytest.mark.parametrize(
+    "redirect_host",
+    [
+        "127.0.0.1",
+        "localhost",
+    ],
+)
+@pytest.mark.parametrize(
+    "redirect_port",
+    [
+        "",
+        ":0",
+        ":12345",
+    ],
+)
+@pytest.mark.parametrize(
+    "dontwait",
+    ["false", "true"],
+)
+@pytest.mark.parametrize("reuse_port", ["true", "false"])
+def test_auth_callback_server_updates_localhost_redirect_uri_port_to_match_socket_port(
+    monkeypatch,
+    socket_host,
+    socket_port,
+    redirect_host,
+    redirect_port,
+    dontwait,
+    reuse_port,
+) -> None:
+    monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", reuse_port)
+    monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT", dontwait)
+    with AuthHttpServer(
+        uri=f"http://{socket_host}{socket_port}/test_request",
+        redirect_uri=f"http://{redirect_host}{redirect_port}/test_request",
+    ) as callback_server:
+        assert callback_server._redirect_uri.port == callback_server.port
+
+
+@pytest.mark.parametrize(
+    "socket_host",
+    [
+        "127.0.0.1",
+        "localhost",
+    ],
+)
+@pytest.mark.parametrize(
+    "socket_port",
+    [
+        "",
+        ":0",
+        ":12345",
+    ],
+)
+@pytest.mark.parametrize(
+    "redirect_host",
+    [
+        "127.0.0.1",
+        "localhost",
+    ],
+)
+@pytest.mark.parametrize(
+    "redirect_port",
+    [
+        54321,
+        54320,
+    ],
+)
+@pytest.mark.parametrize(
+    "dontwait",
+    ["false", "true"],
+)
+@pytest.mark.parametrize("reuse_port", ["true", "false"])
+def test_auth_callback_server_uses_redirect_uri_port_when_specified(
+    monkeypatch,
+    socket_host,
+    socket_port,
+    redirect_host,
+    redirect_port,
+    dontwait,
+    reuse_port,
+) -> None:
+    monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", reuse_port)
+    monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT", dontwait)
+    with AuthHttpServer(
+        uri=f"http://{socket_host}{socket_port}/test_request",
+        redirect_uri=f"http://{redirect_host}:{redirect_port}/test_request",
+    ) as callback_server:
+        assert callback_server.port == redirect_port
+        assert callback_server._redirect_uri.port == redirect_port
+
+
+@pytest.mark.parametrize(
+    "socket_host",
+    [
+        "127.0.0.1",
+        "localhost",
+    ],
+)
+@pytest.mark.parametrize(
+    "socket_port",
+    [
+        "",
+        ":0",
+        ":12345",
+    ],
+)
+@pytest.mark.parametrize(
+    "redirect_port",
+    [
+        "",
+        ":0",
+        ":12345",
+    ],
+)
+@pytest.mark.parametrize(
+    "dontwait",
+    ["false", "true"],
+)
+@pytest.mark.parametrize("reuse_port", ["true", "false"])
+def test_auth_callback_server_does_not_updates_nonlocalhost_redirect_uri_port_to_match_socket_port(
+    monkeypatch, socket_host, socket_port, redirect_port, dontwait, reuse_port
+) -> None:
+    monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", reuse_port)
+    monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT", dontwait)
+    redirect_uri = f"http://not_localhost{redirect_port}/test_request"
+    with AuthHttpServer(
+        uri=f"http://{socket_host}{socket_port}/test_request", redirect_uri=redirect_uri
+    ) as callback_server:
+        assert callback_server.redirect_uri == redirect_uri