separate socket uri and redirect uri

Author: npeshkov

src/snowflake/connector/auth/_http_server.py

@@ -69,9 +69,11 @@ class AuthHttpServer:
     def __init__(
         self,
         uri: str,
+        redirect_uri: str,
         buf_size: int = 16384,
     ) -> None:
         parsed_uri = urllib.parse.urlparse(uri)
+        parsed_redirect = urllib.parse.urlparse(redirect_uri)
         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":
@@ -123,6 +125,24 @@ def __init__(
             query=parsed_uri.query,
             fragment=parsed_uri.fragment,
         )
+        if parsed_redirect.hostname in ("localhost", "127.0.0.1"):
+            logger.debug(
+                f"Redirect URI hostname is {parsed_redirect.hostname}, redirect port {parsed_redirect.port} will be changed to the server port {port}."
+            )
+            self._redirect_uri = urllib.parse.ParseResult(
+                scheme=parsed_redirect.scheme,
+                netloc=parsed_redirect.hostname + ":" + str(port),
+                path=parsed_redirect.path,
+                params=parsed_redirect.params,
+                query=parsed_redirect.query,
+                fragment=parsed_redirect.fragment,
+            )
+        else:
+            self._redirect_uri = parsed_redirect
+
+    @property
+    def redirect_uri(self) -> str:
+        return self._redirect_uri.geturl()
 
     @property
     def url(self) -> str:

src/snowflake/connector/auth/oauth_code.py

@@ -8,6 +8,7 @@
 import hashlib
 import json
 import logging
+import os
 import secrets
 import socket
 import time
@@ -117,7 +118,12 @@ 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(
+            uri=os.environ.get("SNOWFLAKE_OAUTH_SOCKET_ADDRESS", "http://localhost")
+            + ":"
+            + os.environ.get("SNOWFLAKE_OAUTH_SOCKET_PORT", "0"),
+            redirect_uri=self._redirect_uri,
+        ) as callback_server:
             code = self._do_authorization_request(callback_server, conn)
             return self._do_token_request(code, callback_server, conn)
 
@@ -260,7 +266,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(
@@ -314,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"