From 2347911ff336fa2e674914674c6ce9495e470a50 Mon Sep 17 00:00:00 2001
From: David Bochenski <david@goincremental.com>
Date: Sun, 14 Sep 2025 08:41:21 +0000
Subject: [PATCH] MCP OAuth callback + token/consent + Weather MCP policy fixes

- Switch oauth-callback token exchange to client_secret; remove MI assertion
- Align token request scope with configured OAuthScopes
- Fix Cosmos DB partition header: x-ms-documentdb-partitionkey
- Weather MCP policy: validate session against APIM cache (EntraToken-{sessionId})
- Keep encryption session model; no functional change to IV/Key NVs

Validated by pushing policies to APIM and re-running flow: 200 token, Weather MCP accepts bearer token.
---
 .../src/apim-oauth/consent.policy.xml                  |  2 +-
 .../src/apim-oauth/oauth-callback.policy.xml           | 10 ++++++----
 .../src/apim-oauth/token.policy.xml                    |  4 ++--
 .../src/weather/apim-mcp-server/policy.xml             |  9 +++++----
 4 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/labs/mcp-client-authorization/src/apim-oauth/consent.policy.xml b/labs/mcp-client-authorization/src/apim-oauth/consent.policy.xml
index e37d916c..e8507415 100644
--- a/labs/mcp-client-authorization/src/apim-oauth/consent.policy.xml
+++ b/labs/mcp-client-authorization/src/apim-oauth/consent.policy.xml
@@ -242,7 +242,7 @@ __COMMON_STYLES__
                     <set-header name="x-ms-version" exists-action="override">
                         <value>2018-12-31</value>
                     </set-header>
-                    <set-header name="x-ms-partitionkey" exists-action="override">
+                    <set-header name="x-ms-documentdb-partitionkey" exists-action="override">
                         <value>@($"[\"{context.Variables.GetValueOrDefault<string>("client_id")}\"]")</value>
                     </set-header>
                     <set-header name="Authorization" exists-action="override">
diff --git a/labs/mcp-client-authorization/src/apim-oauth/oauth-callback.policy.xml b/labs/mcp-client-authorization/src/apim-oauth/oauth-callback.policy.xml
index 1eb81fe9..9a7b4c7f 100644
--- a/labs/mcp-client-authorization/src/apim-oauth/oauth-callback.policy.xml
+++ b/labs/mcp-client-authorization/src/apim-oauth/oauth-callback.policy.xml
@@ -123,8 +123,8 @@
         <set-variable name="codeChallengeMethod" value="S256" />
         <set-variable name="redirectUri" value="{{OAuthCallbackUri}}" />
         <set-variable name="clientId" value="{{EntraIDClientId}}" />
-        <set-variable name="clientAssertionType" value="@(System.Net.WebUtility.UrlEncode("urn:ietf:params:oauth:client-assertion-type:jwt-bearer"))" />
-        <authentication-managed-identity resource="api://AzureADTokenExchange" client-id="{{EntraIDClientId}}" output-token-variable-name="ficToken"/>
+        <!-- Use confidential client secret for token exchange with Entra ID -->
+        <set-variable name="clientSecret" value="{{EntraIDClientSecret}}" />
          
         <!-- STEP 4: Configure token request to Entra ID -->
         <set-method>POST</set-method>
@@ -132,7 +132,9 @@
             <value>application/x-www-form-urlencoded</value>
         </set-header>
         <set-body>@{
-            return $"client_id={context.Variables.GetValueOrDefault("clientId")}&grant_type=authorization_code&code={context.Variables.GetValueOrDefault("authCode")}&redirect_uri={context.Variables.GetValueOrDefault("redirectUri")}&scope=User.Read&code_verifier={context.Variables.GetValueOrDefault("codeVerifier")}&client_assertion_type={context.Variables.GetValueOrDefault("clientAssertionType")}&client_assertion={context.Variables.GetValueOrDefault("ficToken")}";
+            // Align scope with configured OAuth scopes used during authorize request
+            string scopes = "{{OAuthScopes}}";
+            return $"client_id={context.Variables.GetValueOrDefault("clientId")}&grant_type=authorization_code&code={context.Variables.GetValueOrDefault("authCode")}&redirect_uri={context.Variables.GetValueOrDefault("redirectUri")}&scope={System.Net.WebUtility.UrlEncode(scopes)}&code_verifier={context.Variables.GetValueOrDefault("codeVerifier")}&client_secret={System.Net.WebUtility.UrlEncode((string)context.Variables.GetValueOrDefault("clientSecret"))}";
         }</set-body>
         <rewrite-uri template="/token" />
     </inbound>
@@ -237,4 +239,4 @@
     <on-error>
         <base />
     </on-error>
-</policies>
\ No newline at end of file
+</policies>
diff --git a/labs/mcp-client-authorization/src/apim-oauth/token.policy.xml b/labs/mcp-client-authorization/src/apim-oauth/token.policy.xml
index a3956236..1686b636 100644
--- a/labs/mcp-client-authorization/src/apim-oauth/token.policy.xml
+++ b/labs/mcp-client-authorization/src/apim-oauth/token.policy.xml
@@ -190,7 +190,7 @@
                     <set-header name="x-ms-version" exists-action="override">
                         <value>2018-12-31</value>
                     </set-header>
-                    <set-header name="x-ms-partitionkey" exists-action="override">
+                    <set-header name="x-ms-documentdb-partitionkey" exists-action="override">
                         <value>@($"[\"{context.Variables.GetValueOrDefault<string>("client_id")}\"]")</value>
                     </set-header>
                     <set-header name="Authorization" exists-action="override">
@@ -312,4 +312,4 @@
     <on-error>
         <base />
     </on-error>
-</policies>
\ No newline at end of file
+</policies>
diff --git a/labs/mcp-client-authorization/src/weather/apim-mcp-server/policy.xml b/labs/mcp-client-authorization/src/weather/apim-mcp-server/policy.xml
index 516a63d5..8976288c 100644
--- a/labs/mcp-client-authorization/src/weather/apim-mcp-server/policy.xml
+++ b/labs/mcp-client-authorization/src/weather/apim-mcp-server/policy.xml
@@ -47,9 +47,10 @@
             byte[] decryptedBytes = inBytes.Decrypt("Aes", key, IV);
             return Encoding.UTF8.GetString(decryptedBytes);
          }" />
-        <!-- if decrypted Session Id does not match the session we should throw error-->
+        <!-- Validate session by checking APIM cache for an Entra token tied to this session ID -->
+        <cache-lookup-value key="@($"EntraToken-{context.Variables.GetValueOrDefault("decryptedSessionId", "")}")" variable-name="sessionEntry" />
         <choose>
-            <when condition="@((string)context.Variables["decryptedSessionId"] != "sessionId123")">
+            <when condition="@(string.IsNullOrEmpty((string)context.Variables.GetValueOrDefault("sessionEntry", "")))">
                 <return-response>
                     <set-status code="401" reason="Unauthorized" />
                     <set-header name="WWW-Authenticate" exists-action="override">
@@ -57,7 +58,7 @@
                     </set-header>
                     <set-body>{
                         "error": "unauthorized",
-                        "error_description": "token is not valid"
+                        "error_description": "session not found or expired"
                         }</set-body>
                 </return-response>
             </when>
@@ -72,4 +73,4 @@
     <on-error>
         <base />
     </on-error>
-</policies>
\ No newline at end of file
+</policies>