Handle redirect_uri properly

Author: sam-muncke

src/auth0_server_python/auth_server/server_client.py

@@ -1289,7 +1289,6 @@ async def start_connect_account(
         options: ConnectAccountOptions,
         store_options: dict = None
     ) -> str:
-        
         # Get effective authorization params (merge defaults with provided ones)
         auth_params = dict(self._default_authorization_params)
         if options.authorization_params:
@@ -1298,24 +1297,22 @@ async def start_connect_account(
                 ) if k not in INTERNAL_AUTHORIZE_PARAMS}
             )
 
+        # Use the default redirect_uri if none is specified
+        redirect_uri = options.redirect_uri or self._redirect_uri 
         # Ensure we have a redirect_uri
-        if "redirect_uri" not in auth_params and not self._redirect_uri:
+        if not redirect_uri:
             raise MissingRequiredArgumentError("redirect_uri")
 
-        # Use the default redirect_uri if none is specified
-        if "redirect_uri" not in auth_params and self._redirect_uri:
-            auth_params["redirect_uri"] = self._redirect_uri
-
         # Generate PKCE code verifier and challenge
         code_verifier = PKCE.generate_code_verifier()
         code_challenge = PKCE.generate_code_challenge(code_verifier)
 
         # State parameter to prevent CSRF
         state = PKCE.generate_random_string(32)
-        
+
         connect_request = ConnectAccountRequest(
             connection=options.connection,
-            redirect_uri = options.redirect_uri or auth_params["redirect_uri"],
+            redirect_uri = redirect_uri,
             code_challenge=code_challenge,
             code_challenge_method="S256",
             state=state,
@@ -1335,7 +1332,8 @@ async def start_connect_account(
         transaction_data = TransactionData(
             code_verifier=code_verifier,
             app_state=state,
-            auth_session = connect_response.auth_session
+            auth_session=connect_response.auth_session,
+            redirect_uri=redirect_uri
         )
 
         # Store the transaction data
@@ -1359,18 +1357,17 @@ async def complete_connect_account(
 
         if not transaction_data:
             raise MissingTransactionError()
-        
-        # TODO //do I need to check error in redirect??
-        # TODO //handle no redirect uri??
+
         access_token = await self.get_access_token(
             audience=self._my_account_client.audienceIdentifier,
             scope="create:me:connected_accounts",
             store_options=store_options
         )
+
         request = CompleteConnectAccountRequest(
             auth_session=transaction_data.auth_session,
             connect_code=connect_code,
-            redirect_uri=self._redirect_uri,
+            redirect_uri=transaction_data.redirect_uri,
             code_verifier=transaction_data.code_verifier
         )
 

src/auth0_server_python/auth_types/__init__.py

@@ -88,6 +88,7 @@ class TransactionData(BaseModel):
     code_verifier: str
     app_state: Optional[Any] = None
     auth_session: Optional[str] = None
+    redirect_uri: Optional[str] = None
 
     class Config:
         extra = "allow"  # Allow additional fields not defined in the model

src/auth0_server_python/tests/test_server_client.py

@@ -1268,14 +1268,76 @@ async def test_start_connect_account_calls_connect_and_builds_url(mocker):
     mock_transaction_store = AsyncMock()
     mock_state_store = AsyncMock()
 
+    client = ServerClient(
+        domain="auth0.local",
+        client_id="<client_id>",
+        client_secret="<client_secret>",
+        state_store=mock_state_store,
+        transaction_store=mock_transaction_store,
+        secret="some-secret"
+    )
+    
+    mocker.patch.object(client, "get_access_token", AsyncMock(return_value="<access_token>"))
+    mock_my_account_client = AsyncMock(MyAccountClient)
+    mocker.patch.object(client, "_my_account_client", mock_my_account_client)
+    mock_my_account_client.connect_account.return_value = ConnectAccountResponse(
+        auth_session="<auth_session>",
+        connect_uri="http://auth0.local/connected_accounts/connect",
+        connect_params=ConnectParams(
+            ticket="ticket123",
+        ),
+        expires_in=300
+    )
+
+    mocker.patch.object(PKCE, "generate_random_string", return_value="<state>")
+    mocker.patch.object(PKCE, "generate_code_verifier", return_value="<code_verifier>")
+    mocker.patch.object(PKCE, "generate_code_challenge", return_value="<code_challenge>")
+
+    # Act
+    url = await client.start_connect_account(
+        options=ConnectAccountOptions(
+            connection="<connection>",
+            redirect_uri="/test_redirect_uri"
+        )
+    )
+
+    # Assert
+    assert url == "http://auth0.local/connected_accounts/connect?ticket=ticket123"
+    mock_my_account_client.connect_account.assert_awaited_with(
+        access_token="<access_token>",
+        request=ConnectAccountRequest(
+            connection="<connection>",
+            redirect_uri="/test_redirect_uri",
+            code_challenge_method="S256",
+            code_challenge="<code_challenge>",
+            state= "<state>"
+        )
+    )
+    mock_transaction_store.set.assert_awaited_with(
+        "_a0_tx:<state>", 
+        TransactionData(
+            code_verifier="<code_verifier>",
+            app_state="<state>",
+            auth_session="<auth_session>",
+            redirect_uri="/test_redirect_uri"
+        ),
+        options=ANY
+    )
+
+@pytest.mark.asyncio
+async def test_start_connect_account_default_redirect_uri(mocker):
+    # Setup
+    mock_transaction_store = AsyncMock()
+    mock_state_store = AsyncMock()
+
     client = ServerClient(
         domain="auth0.local",
         client_id="<client_id>",
         client_secret="<client_secret>",
         state_store=mock_state_store,
         transaction_store=mock_transaction_store,
         secret="some-secret",
-        redirect_uri="/test_redirect_uri"
+        redirect_uri="/default_redirect_uri"
     )
 
     mocker.patch.object(client, "get_access_token", AsyncMock(return_value="<access_token>"))
@@ -1298,7 +1360,7 @@ async def test_start_connect_account_calls_connect_and_builds_url(mocker):
     url = await client.start_connect_account(
         options=ConnectAccountOptions(
             connection="<connection>"
-        ),
+        )
     )
 
     # Assert
@@ -1307,7 +1369,7 @@ async def test_start_connect_account_calls_connect_and_builds_url(mocker):
         access_token="<access_token>",
         request=ConnectAccountRequest(
             connection="<connection>",
-            redirect_uri="/test_redirect_uri",
+            redirect_uri="/default_redirect_uri",
             code_challenge_method="S256",
             code_challenge="<code_challenge>",
             state= "<state>"
@@ -1319,10 +1381,37 @@ async def test_start_connect_account_calls_connect_and_builds_url(mocker):
             code_verifier="<code_verifier>",
             app_state="<state>",
             auth_session="<auth_session>",
+            redirect_uri="/default_redirect_uri"
         ),
         options=ANY
     )
 
+@pytest.mark.asyncio
+async def test_start_connect_account_no_redirect_uri(mocker):
+    # Setup
+    mock_transaction_store = AsyncMock()
+    mock_state_store = AsyncMock()
+
+    client = ServerClient(
+        domain="auth0.local",
+        client_id="<client_id>",
+        client_secret="<client_secret>",
+        state_store=mock_state_store,
+        transaction_store=mock_transaction_store,
+        secret="some-secret"
+    )
+
+    # Act
+    with pytest.raises(MissingRequiredArgumentError) as exc:
+        await client.start_connect_account(
+        options=ConnectAccountOptions(
+            connection="<connection>"
+        )
+    )
+
+    # Assert
+    assert "redirect_uri" in str(exc.value)
+
 @pytest.mark.asyncio
 async def test_complete_connect_account_calls_complete(mocker):
     # Setup
@@ -1347,6 +1436,7 @@ async def test_complete_connect_account_calls_complete(mocker):
         code_verifier="<code_verifier>",
         app_state="<state>",
         auth_session="<auth_session>",
+        redirect_uri="/test_redirect_uri"
     )
 
     # Act
@@ -1366,7 +1456,6 @@ async def test_complete_connect_account_calls_complete(mocker):
         )
     )
 
-
 @pytest.mark.asyncio
 async def test_complete_connect_account_no_transactions(mocker):
     # Setup