Improve wrapped key loader

changelog.txt

@@ -1,5 +1,6 @@
 vNext
 ----------
+- [PATCH] Fix caching of secret key and add retries for InvalidKeyException during unwrap (#2659)
 - [MINOR] Replace AbstractSecretKeyLoader with ISecretKeyProvider (#2666)
 - [MINOR] Update IP phone app teams signature constants to use SHA-512 format (#2700)
 - [MINOR] Updating handling of ssl error received in Android WebView's onReceivedSslError callback (#2691)

common/src/androidTest/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProviderTest.java

@@ -22,6 +22,8 @@
 // THE SOFTWARE.
 package com.microsoft.identity.common.crypto;
 
+import static com.microsoft.identity.common.java.exception.ClientException.INVALID_KEY;
+
 import android.content.Context;
 import android.security.KeyPairGeneratorSpec;
 
@@ -47,8 +49,6 @@
 import javax.crypto.SecretKey;
 import javax.security.auth.x500.X500Principal;
 
-import static com.microsoft.identity.common.java.exception.ClientException.INVALID_KEY;
-
 public class AndroidWrappedKeyProviderTest {
 
     final Context context = ApplicationProvider.getApplicationContext();
@@ -144,7 +144,7 @@ public void testLoadKey() throws ClientException {
         final AndroidWrappedKeyProvider keyProvider = new AndroidWrappedKeyProvider(MOCK_KEY_ALIAS, MOCK_KEY_FILE_PATH, context);
         final SecretKey secretKey = keyProvider.getKey();
 
-        final SecretKey key = keyProvider.getKeyCache().getData();
+        final SecretKey key = keyProvider.getKeyFromCache();
         Assert.assertNotNull(key);
         Assert.assertEquals(AES_ALGORITHM, secretKey.getAlgorithm());
         Assert.assertArrayEquals(secretKey.getEncoded(), key.getEncoded());
@@ -228,7 +228,7 @@ public void testPerf_NoCachedKey() throws ClientException {
 
         long timeStartLoopNotCached = System.nanoTime();
         for (int i = 0; i < 100; i++) {
-            keyProvider.getKeyCache().clear();
+            keyProvider.clearKeyFromCache();
             keyProvider.getKey();
         }
         long timeFinishLoopNotCached = System.nanoTime();
@@ -246,7 +246,7 @@ public void testLoadDeletedKeyStoreKey() throws ClientException {
         AndroidKeyStoreUtil.deleteKey(MOCK_KEY_ALIAS);
 
         // Cached key also be wiped.
-        final SecretKey key = keyProvider.getKeyCache().getData();
+        final SecretKey key = keyProvider.getKeyFromCache();
         Assert.assertNull(key);
     }
 
@@ -257,15 +257,15 @@ public void testLoadDeletedKeyFile() throws ClientException {
         FileUtil.deleteFile(getKeyFile());
 
         // Cached key also be wiped.
-        final SecretKey key = keyProvider.getKeyCache().getData();
+        final SecretKey key = keyProvider.getKeyFromCache();
         Assert.assertNull(key);
     }
 
     private AndroidWrappedKeyProvider initkeyProviderWithKeyEntry() throws ClientException {
         final AndroidWrappedKeyProvider keyProvider = new AndroidWrappedKeyProvider(MOCK_KEY_ALIAS, MOCK_KEY_FILE_PATH, context);
         final SecretKey key = keyProvider.getKey();
         Assert.assertNotNull(key);
-        Assert.assertNotNull(keyProvider.getKeyCache().getData());
+        Assert.assertNotNull(keyProvider.getKeyFromCache());
         return keyProvider;
     }
 }

common/src/androidTest/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtilTest.java

@@ -0,0 +1,117 @@
+// Copyright (c) Microsoft Corporation.
+// All rights reserved.
+//
+// This code is licensed under the MIT License.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files(the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and / or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions :
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package com.microsoft.identity.common.internal.util;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import com.microsoft.identity.common.java.exception.ClientException;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+
+import java.security.KeyPair;
+import java.security.KeyPairGenerator;
+import java.security.NoSuchAlgorithmException;
+
+import javax.crypto.Cipher;
+import javax.crypto.SecretKey;
+import javax.crypto.spec.SecretKeySpec;
+
+@RunWith(AndroidJUnit4.class)
+public class AndroidKeyStoreUtilTest {
+    private static final String TEST_WRAP_ALGORITHM = "RSA/ECB/PKCS1Padding";
+    private static final String TEST_KEY_ALGORITHM = "AES";
+    private static final byte[] TEST_KEY_BYTES = new byte[]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16};
+    private static final byte[] TEST_WRAPPED_KEY_BYTES = new byte[]{16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1};
+
+    private KeyPair mKeyPair;
+    private SecretKey mSecretKey;
+
+    @Before
+    public void setUp() throws Exception {
+        // Create a real KeyPair for testing
+        KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
+        keyGen.initialize(2048);
+        mKeyPair = keyGen.generateKeyPair();
+
+        // Create a test SecretKey
+        mSecretKey = new SecretKeySpec(TEST_KEY_BYTES, TEST_KEY_ALGORITHM);
+    }
+
+    @Test
+    public void testUnwrap_Success() throws Exception {
+        // Mock Cipher for successful unwrap
+        try (MockedStatic<Cipher> mockedCipher = Mockito.mockStatic(Cipher.class)) {
+            Cipher mockCipher = mock(Cipher.class);
+            when(mockCipher.unwrap(any(byte[].class), any(String.class), any(int.class)))
+                    .thenReturn(mSecretKey);
+            mockedCipher.when(() -> Cipher.getInstance(TEST_WRAP_ALGORITHM))
+                    .thenReturn(mockCipher);
+
+            // Test successful unwrap
+            SecretKey result = AndroidKeyStoreUtil.unwrap(
+                    TEST_WRAPPED_KEY_BYTES,
+                    TEST_KEY_ALGORITHM,
+                    mKeyPair,
+                    TEST_WRAP_ALGORITHM
+            );
+
+            assertNotNull(result);
+            assertEquals(mSecretKey, result);
+            verify(mockCipher, times(1)).init(Cipher.UNWRAP_MODE, mKeyPair.getPrivate());
+        }
+    }
+
+    @Test
+    public void testUnwrap_NoSuchAlgorithm() throws Exception {
+        try (MockedStatic<Cipher> mockedCipher = Mockito.mockStatic(Cipher.class)) {
+            mockedCipher.when(() -> Cipher.getInstance(TEST_WRAP_ALGORITHM))
+                    .thenThrow(new NoSuchAlgorithmException("Test failure"));
+
+            try {
+                AndroidKeyStoreUtil.unwrap(
+                        TEST_WRAPPED_KEY_BYTES,
+                        TEST_KEY_ALGORITHM,
+                        mKeyPair,
+                        TEST_WRAP_ALGORITHM
+                );
+                fail("Should have thrown ClientException");
+            } catch (ClientException e) {
+                assertEquals("NO_SUCH_ALGORITHM", e.getErrorCode());
+            }
+        }
+    }
+}

common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java

@@ -44,7 +44,6 @@
 import com.microsoft.identity.common.java.opentelemetry.OTelUtility;
 import com.microsoft.identity.common.java.opentelemetry.SpanExtension;
 import com.microsoft.identity.common.java.opentelemetry.SpanName;
-import com.microsoft.identity.common.java.util.CachedData;
 import com.microsoft.identity.common.java.util.FileUtil;
 import com.microsoft.identity.common.java.util.StringUtil;
 import com.microsoft.identity.common.logging.Logger;
@@ -57,6 +56,8 @@
 import java.security.spec.AlgorithmParameterSpec;
 import java.util.Calendar;
 import java.util.Locale;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 
 import javax.crypto.SecretKey;
 import javax.security.auth.x500.X500Principal;
@@ -109,6 +110,10 @@ public class AndroidWrappedKeyProvider implements ISecretKeyProvider {
     // Exposed for testing only.
     /* package */ static final int KEY_FILE_SIZE = 1024;
 
+    /**
+     * SecretKey cache. Maps wrapped secret key file path to the SecretKey.
+     */
+    private static final ConcurrentMap<String, SecretKey> sKeyCacheMap = new ConcurrentHashMap<>();
     private final Context mContext;
 
     /**
@@ -122,21 +127,15 @@ public class AndroidWrappedKeyProvider implements ISecretKeyProvider {
      */
     private final String mFilePath;
 
-    private final CachedData<SecretKey> mKeyCache = new CachedData<SecretKey>() {
-        @Override
-        public SecretKey getData() {
-            if (!sSkipKeyInvalidationCheck &&
-                    (!AndroidKeyStoreUtil.canLoadKey(mAlias) || !getKeyFile().exists())) {
-                this.clear();
-            }
-            return super.getData();
-        }
-    };
+    // Exposed for testing only.
+    @Nullable
+    /* package */ SecretKey getKeyFromCache() {
+        return sKeyCacheMap.get(mFilePath);
+    }
 
     // Exposed for testing only.
-    @NonNull
-    /* package */ CachedData<SecretKey> getKeyCache() {
-        return mKeyCache;
+    /* package */ void clearKeyFromCache() {
+        sKeyCacheMap.remove(mFilePath);
     }
 
     /**
@@ -173,18 +172,27 @@ public String getKeyTypeIdentifier() {
     @Override
     @NonNull
     public synchronized SecretKey getKey() throws ClientException {
-        SecretKey key = mKeyCache.getData();
+        final String methodTag = TAG + ":getKey";
+        if (!sSkipKeyInvalidationCheck &&
+                (!AndroidKeyStoreUtil.canLoadKey(mAlias) || !this.getKeyFile().exists())) {
+            sKeyCacheMap.remove(mFilePath);
+        }
 
-        if (key == null) {
-            key = readSecretKeyFromStorage();
+        SecretKey key = sKeyCacheMap.get(mFilePath);
+        if (key != null) {
+            return key;
         }
 
+        Logger.info(methodTag, "Key not in cache is empty, loading key from storage");
+        key = readSecretKeyFromStorage();
+
         // If key doesn't exist, generate a new one.
         if (key == null) {
+            Logger.info(methodTag, "Key does not exist in storage, generating a new key");
             key = generateRandomKey();
         }
 
-        mKeyCache.setData(key);
+        sKeyCacheMap.put(mFilePath, key);
         return key;
     }
 
@@ -223,7 +231,7 @@ protected SecretKey generateRandomKey() throws ClientException {
                 // Do not delete the KeyStoreKeyPair even if the key file is empty. This caused credential cache
                 // to be deleted in Office because of sharedUserId allowing keystore to be shared amongst apps.
                 FileUtil.deleteFile(getKeyFile());
-                mKeyCache.clear();
+                sKeyCacheMap.remove(mFilePath);
                 return null;
             }
 
@@ -427,7 +435,7 @@ private void recordKeyGenerationTime(long keypairGenStartTime) {
     public void deleteSecretKeyFromStorage() throws ClientException {
         AndroidKeyStoreUtil.deleteKey(mAlias);
         FileUtil.deleteFile(getKeyFile());
-        mKeyCache.clear();
+        sKeyCacheMap.remove(mFilePath);
     }
 
     /**

common/src/main/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtil.java

@@ -87,20 +87,16 @@ public class AndroidKeyStoreUtil {
     private AndroidKeyStoreUtil() {
     }
 
-    private static KeyStore mKeyStore;
-
     private static final LongCounter sFailedAndroidKeyStoreUnwrapOperationCount = OTelUtility.createLongCounter(
             "failed_keystore_key_unwrap_operation_count",
             "Number of failed Android KeyStore unwrap operations"
     );
 
     private static synchronized KeyStore getKeyStore()
             throws KeyStoreException, CertificateException, NoSuchAlgorithmException, IOException {
-        if (mKeyStore == null){
-            mKeyStore = KeyStore.getInstance(ANDROID_KEY_STORE_TYPE);
-            mKeyStore.load(null);
-        }
-        return mKeyStore;
+        final KeyStore keyStore = KeyStore.getInstance(ANDROID_KEY_STORE_TYPE);
+        keyStore.load(null);
+        return keyStore;
     }
 
     /**
@@ -307,8 +303,7 @@ public static synchronized void deleteKey(
         final Throwable exception;
         final String errCode;
         try {
-            final KeyStore keyStore = KeyStore.getInstance(ANDROID_KEY_STORE_TYPE);
-            keyStore.load(null);
+            final KeyStore keyStore = getKeyStore();
             keyStore.deleteEntry(aliasOfKeyToDelete);
             return;
         } catch (final KeyStoreException e) {