use redirect port when specified

Author: npeshkov

src/snowflake/connector/auth/_http_server.py

@@ -84,7 +84,10 @@ def __init__(
             else:
                 self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
 
-        port = parsed_uri.port or 0
+        if parsed_redirect.hostname in ("localhost", "127.0.0.1"):
+            port = parsed_redirect.port or 0
+        else:
+            port = parsed_uri.port or 0
         for attempt in range(1, self.DEFAULT_MAX_ATTEMPTS + 1):
             try:
                 self._socket.bind(
@@ -125,9 +128,12 @@ def __init__(
             query=parsed_uri.query,
             fragment=parsed_uri.fragment,
         )
-        if parsed_redirect.hostname in ("localhost", "127.0.0.1"):
+        if (
+            parsed_redirect.hostname in ("localhost", "127.0.0.1")
+            and port != parsed_redirect.port
+        ):
             logger.debug(
-                f"Redirect URI hostname is {parsed_redirect.hostname}, redirect port {parsed_redirect.port} will be changed to the server port {port}."
+                f"Updating redirect port {parsed_redirect.port} to match the server port {port}."
             )
             self._redirect_uri = urllib.parse.ParseResult(
                 scheme=parsed_redirect.scheme,

test/unit/test_auth_callback_server.py

@@ -122,6 +122,59 @@ def test_auth_callback_server_updates_localhost_redirect_uri_port_to_match_socke
         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",
     [