Fix JwtSecurityToken Missing Mapping When Creating a Token. (#2578)

* Added test for debugging. * Commit for draft PR * Fix with AppContext switch * Added test for debugging. * Commit for draft PR * Fix with AppContext switch * Clean-up commit * One more clean-up * Fixes JsonWebTokenHandler as well * Disable parallelization on test due to AppContext being a static dict * Move AppContext switch test to separate file * Update test/System.IdentityModel.Tokens.Jwt.Tests/JwtSecurityTokenHandlerNonParallelRunTests.cs Co-authored-by: kellyyangsong <69649063+kellyyangsong@users.noreply.github.com> * Update test/System.IdentityModel.Tokens.Jwt.Tests/JwtSecurityTokenHandlerTests.WithContextSwitches.cs Co-authored-by: kellyyangsong <69649063+kellyyangsong@users.noreply.github.com> * Fix blank space * Address PR feedback --------- Co-authored-by: Franco Fung <francofung@microsoft.com> Co-authored-by: kellyyangsong <69649063+kellyyangsong@users.noreply.github.com>
AzureAD · May 20, 2024 · 0e5648b · 0e5648b
1 parent ad64ace
commit 0e5648b
Show file tree

Hide file tree

Showing 3 changed files with 80 additions and 2 deletions.
diff --git a/src/Microsoft.IdentityModel.Tokens/X509EncryptingCredentials.cs b/src/Microsoft.IdentityModel.Tokens/X509EncryptingCredentials.cs
@@ -11,6 +11,8 @@ namespace Microsoft.IdentityModel.Tokens
     /// </summary>
     public class X509EncryptingCredentials : EncryptingCredentials
     {
+        internal const string _useShortNameForRsaOaepKey = "Switch.Microsoft.IdentityModel.UseShortNameForRsaOaepKey";
+
         /// <summary>
         /// Designed to construct <see cref="EncryptingCredentials"/> based on a x509 certificate.
         /// </summary>
@@ -21,7 +23,7 @@ public class X509EncryptingCredentials : EncryptingCredentials
         /// </remarks>
         /// <exception cref="ArgumentNullException">if 'certificate' is null.</exception>
         public X509EncryptingCredentials(X509Certificate2 certificate)
-            : this(certificate, SecurityAlgorithms.DefaultAsymmetricKeyWrapAlgorithm, SecurityAlgorithms.DefaultSymmetricEncryptionAlgorithm)
+            : this(certificate, GetEncryptionAlgorithm(), SecurityAlgorithms.DefaultSymmetricEncryptionAlgorithm)
         {
         }
 
@@ -48,5 +50,15 @@ public X509Certificate2 Certificate
             get;
             private set;
         }
+
+        private static string GetEncryptionAlgorithm()
+        {
+            return ShouldUseShortNameForRsaOaepKey() ? SecurityAlgorithms.RsaOAEP : SecurityAlgorithms.DefaultAsymmetricKeyWrapAlgorithm;
+        }
+
+        private static bool ShouldUseShortNameForRsaOaepKey()
+        {
+            return AppContext.TryGetSwitch(_useShortNameForRsaOaepKey, out var useKeyWrap) && useKeyWrap;
+        }
     }
 }
diff --git a/src/System.IdentityModel.Tokens.Jwt/JwtSecurityTokenHandler.cs b/src/System.IdentityModel.Tokens.Jwt/JwtSecurityTokenHandler.cs
@@ -35,7 +35,7 @@ public class JwtSecurityTokenHandler : SecurityTokenHandler
         private Dictionary<string, string> _outboundAlgorithmMap = null;
         private static string _shortClaimType = _namespace + "/ShortTypeName";
         private bool _mapInboundClaims = DefaultMapInboundClaims;
-
+       
         /// <summary>
         /// Default claim type mapping for inbound claims.
         /// </summary>

diff --git a/...System.IdentityModel.Tokens.Jwt.Tests/JwtSecurityTokenHandlerTests.WithContextSwitches.cs b/...System.IdentityModel.Tokens.Jwt.Tests/JwtSecurityTokenHandlerTests.WithContextSwitches.cs
@@ -0,0 +1,66 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using Microsoft.IdentityModel.JsonWebTokens;
+using Microsoft.IdentityModel.TestUtils;
+using Microsoft.IdentityModel.Tokens;
+using Xunit;
+
+namespace System.IdentityModel.Tokens.Jwt.Tests
+{
+    [CollectionDefinition("JwtSecurityTokenHandlerTestsWithContextSwitches", DisableParallelization = true)]
+    public class JwtSecurityTokenHandlerTestsWithContextSwitches
+    {
+        [Theory]
+        [InlineData(SecurityAlgorithms.RsaOAEP, true)]
+        [InlineData(SecurityAlgorithms.RsaOaepKeyWrap, false)]
+        public void JwtSecurityTokenHandler_CreateToken_AddShortFormMappingForRsaOAEP(string algorithm, bool useShortNameForRsaOaepKey)
+        {
+            AppContext.SetSwitch(X509EncryptingCredentials._useShortNameForRsaOaepKey, useShortNameForRsaOaepKey);
+            var encryptingCredentials = new X509EncryptingCredentials(Default.Certificate);
+            JwtSecurityTokenHandler tokenHandler = new JwtSecurityTokenHandler();
+
+            JwtSecurityToken token = CreateJwtSecurityToken(tokenHandler, encryptingCredentials);
+
+            Assert.Equal(token.Header.Alg, algorithm);
+
+            AppContext.SetSwitch(X509EncryptingCredentials._useShortNameForRsaOaepKey, false);
+        }
+
+        [Theory]
+        [InlineData(SecurityAlgorithms.RsaOAEP, true)]
+        [InlineData(SecurityAlgorithms.RsaOaepKeyWrap, false)]
+        public void JsonWebTokenHandler_CreateToken_AddShortFormMappingForRsaOAEP(string algorithm, bool useShortNameForRsaOaepKey)
+        {
+            AppContext.SetSwitch(X509EncryptingCredentials._useShortNameForRsaOaepKey, useShortNameForRsaOaepKey);
+            var encryptingCredentials = new X509EncryptingCredentials(Default.Certificate);
+            JsonWebTokenHandler tokenHandler = new JsonWebTokenHandler();
+
+            JsonWebToken jsonToken = new JsonWebToken(CreateJwtSecurityTokenAsString(tokenHandler, encryptingCredentials));
+
+            Assert.Equal(jsonToken.Alg, algorithm);
+
+            AppContext.SetSwitch(X509EncryptingCredentials._useShortNameForRsaOaepKey, false);
+        }
+
+        private JwtSecurityToken CreateJwtSecurityToken(JwtSecurityTokenHandler tokenHandler, X509EncryptingCredentials encryptingCredentials)
+        {
+            return tokenHandler.CreateJwtSecurityToken(CreateTokenDescriptor(encryptingCredentials));
+        }
+
+        private string CreateJwtSecurityTokenAsString(JsonWebTokenHandler tokenHandler, X509EncryptingCredentials encryptingCredentials)
+        {
+            return tokenHandler.CreateToken(CreateTokenDescriptor(encryptingCredentials));
+        }
+
+        private SecurityTokenDescriptor CreateTokenDescriptor(X509EncryptingCredentials encryptingCredentials)
+        {
+            return new SecurityTokenDescriptor
+            {
+                Issuer = Default.Issuer,
+                SigningCredentials = Default.AsymmetricSigningCredentials,
+                EncryptingCredentials = encryptingCredentials,
+            };
+        }
+    }
+}