From ec07ebb93e62bed863911a39b17a239bfb5ca80c Mon Sep 17 00:00:00 2001 From: Mohit Date: Thu, 29 May 2025 16:45:33 -0700 Subject: [PATCH 1/9] improve wrapped key loader --- changelog.txt | 1 + .../crypto/AndroidWrappedKeyLoaderTest.java | 117 +++++++++++- .../util/AndroidKeyStoreUtilTest.java | 175 ++++++++++++++++++ .../crypto/AndroidWrappedKeyLoader.java | 48 +++-- .../internal/util/AndroidKeyStoreUtil.java | 78 +++++--- .../java/crypto/StorageEncryptionManager.java | 3 +- 6 files changed, 370 insertions(+), 52 deletions(-) create mode 100644 common/src/androidTest/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtilTest.java diff --git a/changelog.txt b/changelog.txt index 9735a5d1a1..62e6576590 100644 --- a/changelog.txt +++ b/changelog.txt @@ -1,5 +1,6 @@ vNext ---------- +- [PATCH] Fix caching of secret key and add retries for InvalidKeyException during unwrap - [MINOR] Add Get AAD device ID API for resource account (#2644) - [MINOR] Expose a JavaScript API in brokered Webviews to facilitate Improved Same Device NumberMatch (#2617) - [MINOR] Add API for resource account provisioning (API only) (#2640) diff --git a/common/src/androidTest/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoaderTest.java b/common/src/androidTest/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoaderTest.java index 177e011c89..7d59cbb5d5 100644 --- a/common/src/androidTest/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoaderTest.java +++ b/common/src/androidTest/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoaderTest.java @@ -31,6 +31,7 @@ import com.microsoft.identity.common.internal.util.AndroidKeyStoreUtil; import com.microsoft.identity.common.java.crypto.key.AES256KeyLoader; import com.microsoft.identity.common.java.exception.ClientException; +import com.microsoft.identity.common.java.util.CachedData; import com.microsoft.identity.common.java.util.FileUtil; import org.junit.Assert; @@ -43,7 +44,14 @@ import java.security.KeyPair; import java.security.spec.AlgorithmParameterSpec; import java.util.Arrays; +import java.util.Collections; import java.util.Date; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.ArrayList; import javax.crypto.SecretKey; import javax.security.auth.x500.X500Principal; @@ -108,6 +116,7 @@ private AlgorithmParameterSpec getMockKeyPairGeneratorSpec(final String alias) { .build(); } + @Test public void testGenerateKey() throws ClientException { final AndroidWrappedKeyLoader keyLoader = new AndroidWrappedKeyLoader(MOCK_KEY_ALIAS, MOCK_KEY_FILE_PATH, context); @@ -142,8 +151,9 @@ public void testLoadKey() throws ClientException { final AndroidWrappedKeyLoader keyLoader = new AndroidWrappedKeyLoader(MOCK_KEY_ALIAS, MOCK_KEY_FILE_PATH, context); final SecretKey secretKey = keyLoader.getKey(); - - final SecretKey key = keyLoader.getKeyCache().getData(); + final CachedData cachedKey = keyLoader.getKeyCache(); + Assert.assertNotNull(cachedKey); + final SecretKey key = cachedKey.getData(); Assert.assertNotNull(key); Assert.assertEquals(AES256KeyLoader.AES_ALGORITHM, secretKey.getAlgorithm()); Assert.assertArrayEquals(secretKey.getEncoded(), key.getEncoded()); @@ -151,7 +161,7 @@ public void testLoadKey() throws ClientException { } @Test - public void testLoadKeyFromCorruptedFile_TruncatedExisingKey() throws ClientException { + public void testLoadKeyFromCorruptedFile_TruncatedExistingKey() throws ClientException { // Create a new Keystore-wrapped key. final AndroidWrappedKeyLoader keyLoader = new AndroidWrappedKeyLoader(MOCK_KEY_ALIAS, MOCK_KEY_FILE_PATH, context); keyLoader.generateRandomKey(); @@ -166,7 +176,7 @@ public void testLoadKeyFromCorruptedFile_TruncatedExisingKey() throws ClientExce try{ keyLoader.readSecretKeyFromStorage(); Assert.fail(); - } catch (ClientException e){ + } catch (final ClientException e){ Assert.assertEquals(INVALID_KEY, e.getErrorCode()); } @@ -227,7 +237,7 @@ public void testPerf_NoCachedKey() throws ClientException { long timeStartLoopNotCached = System.nanoTime(); for (int i = 0; i < 100; i++) { - keyLoader.getKeyCache().clear(); + keyLoader.clearKeyCache(); keyLoader.getKey(); } long timeFinishLoopNotCached = System.nanoTime(); @@ -245,7 +255,7 @@ public void testLoadDeletedKeyStoreKey() throws ClientException { AndroidKeyStoreUtil.deleteKey(MOCK_KEY_ALIAS); // Cached key also be wiped. - final SecretKey key = keyLoader.getKeyCache().getData(); + final CachedData key = keyLoader.getKeyCache(); Assert.assertNull(key); } @@ -256,15 +266,106 @@ public void testLoadDeletedKeyFile() throws ClientException { FileUtil.deleteFile(getKeyFile()); // Cached key also be wiped. - final SecretKey key = keyLoader.getKeyCache().getData(); + final CachedData key = keyLoader.getKeyCache(); Assert.assertNull(key); } + @Test + public void testGetKey_CacheWorks() throws ClientException { + final AndroidWrappedKeyLoader keyLoader = initKeyLoaderWithKeyEntry(); + + // First call should read from storage + final SecretKey key1 = keyLoader.getKey(); + // Second call should return cached value + final SecretKey key2 = keyLoader.getKey(); + + // Same key should be returned (cached) + Assert.assertSame(key1, key2); + } + + @Test + public void testGetKey_DifferentAliasGetsDifferentCache() throws ClientException { + final AndroidWrappedKeyLoader keyLoader1 = new AndroidWrappedKeyLoader(MOCK_KEY_ALIAS, MOCK_KEY_FILE_PATH, context); + final AndroidWrappedKeyLoader keyLoader2 = new AndroidWrappedKeyLoader(MOCK_KEY_ALIAS + "2", MOCK_KEY_FILE_PATH + "2", context); + + // Generate keys for both loaders + final SecretKey key1 = keyLoader1.getKey(); + final SecretKey key2 = keyLoader2.getKey(); + + // Keys should be different since they have different aliases + Assert.assertNotSame(key1, key2); + Assert.assertFalse(Arrays.equals(key1.getEncoded(), key2.getEncoded())); + } + + @Test + public void testGetKey_SameAliasSharesCache() throws ClientException { + final AndroidWrappedKeyLoader keyLoader1 = new AndroidWrappedKeyLoader(MOCK_KEY_ALIAS, MOCK_KEY_FILE_PATH, context); + final AndroidWrappedKeyLoader keyLoader2 = new AndroidWrappedKeyLoader(MOCK_KEY_ALIAS, MOCK_KEY_FILE_PATH + "different", context); + + // Get key from first loader + final SecretKey key1 = keyLoader1.getKey(); + // Second loader should get same key from cache despite different file path + final SecretKey key2 = keyLoader2.getKey(); + + // Keys should be the same since they share the same alias + Assert.assertSame(key1, key2); + } + + @Test + public void testGetKey_ConcurrentAccess() throws Exception { + final AndroidWrappedKeyLoader keyLoader = new AndroidWrappedKeyLoader(MOCK_KEY_ALIAS, MOCK_KEY_FILE_PATH, context); + final int threadCount = 10; + final CountDownLatch startLatch = new CountDownLatch(1); + final CountDownLatch finishLatch = new CountDownLatch(threadCount); + final Set keys = Collections.synchronizedSet(new HashSet<>()); + final List exceptions = Collections.synchronizedList(new ArrayList<>()); + + // Create multiple threads that will try to get the key simultaneously + for (int i = 0; i < threadCount; i++) { + new Thread(() -> { + try { + startLatch.await(); // Wait for all threads to be ready + keys.add(keyLoader.getKey()); + } catch (Exception e) { + exceptions.add(e); + } finally { + finishLatch.countDown(); + } + }).start(); + } + + startLatch.countDown(); // Start all threads + finishLatch.await(2, TimeUnit.SECONDS); // Wait for all threads to finish + + // Check no exceptions occurred + Assert.assertTrue("Exceptions occurred: " + exceptions, exceptions.isEmpty()); + // All threads should get the same key instance + Assert.assertEquals("All threads should get the same key", 1, keys.size()); + } + + @Test + public void testGetKey_DeletedKeyRegenerated() throws ClientException { + final AndroidWrappedKeyLoader keyLoader = initKeyLoaderWithKeyEntry(); + + // Get initial key + final SecretKey initialKey = keyLoader.getKey(); + + // Delete the key + keyLoader.deleteSecretKeyFromStorage(); + + // Get key again - should generate a new one + final SecretKey newKey = keyLoader.getKey(); + + // Keys should be different + Assert.assertNotSame(initialKey, newKey); + Assert.assertFalse(Arrays.equals(initialKey.getEncoded(), newKey.getEncoded())); + } + private AndroidWrappedKeyLoader initKeyLoaderWithKeyEntry() throws ClientException { final AndroidWrappedKeyLoader keyLoader = new AndroidWrappedKeyLoader(MOCK_KEY_ALIAS, MOCK_KEY_FILE_PATH, context); final SecretKey key = keyLoader.getKey(); Assert.assertNotNull(key); - Assert.assertNotNull(keyLoader.getKeyCache().getData()); + Assert.assertNotNull(keyLoader.getKeyCache()); return keyLoader; } } diff --git a/common/src/androidTest/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtilTest.java b/common/src/androidTest/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtilTest.java new file mode 100644 index 0000000000..02ed1fcb6c --- /dev/null +++ b/common/src/androidTest/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtilTest.java @@ -0,0 +1,175 @@ +// 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.when; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +import android.security.keystore.KeyGenParameterSpec; +import android.security.keystore.KeyProperties; + +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.InvalidKeyException; +import java.security.KeyPair; +import java.security.KeyPairGenerator; +import java.security.KeyStore; +import java.security.NoSuchAlgorithmException; +import java.security.PrivateKey; +import java.security.PublicKey; + +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 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_WithRetries() throws Exception { + // Mock Cipher to fail twice with InvalidKeyException then succeed + try (MockedStatic mockedCipher = Mockito.mockStatic(Cipher.class)) { + Cipher mockCipher = mock(Cipher.class); + when(mockCipher.unwrap(any(byte[].class), any(String.class), any(int.class))) + .thenThrow(new InvalidKeyException("Test failure 1")) + .thenThrow(new InvalidKeyException("Test failure 2")) + .thenReturn(mSecretKey); + mockedCipher.when(() -> Cipher.getInstance(TEST_WRAP_ALGORITHM)) + .thenReturn(mockCipher); + + // Test unwrap with retries + SecretKey result = AndroidKeyStoreUtil.unwrap( + TEST_WRAPPED_KEY_BYTES, + TEST_KEY_ALGORITHM, + mKeyPair, + TEST_WRAP_ALGORITHM + ); + + assertNotNull(result); + assertEquals(mSecretKey, result); + verify(mockCipher, times(3)).init(Cipher.UNWRAP_MODE, mKeyPair.getPrivate()); + } + } + + @Test + public void testUnwrap_MaxRetriesExceeded() throws Exception { + // Mock Cipher to always fail with InvalidKeyException + try (MockedStatic mockedCipher = Mockito.mockStatic(Cipher.class)) { + Cipher mockCipher = mock(Cipher.class); + when(mockCipher.unwrap(any(byte[].class), any(String.class), any(int.class))) + .thenThrow(new InvalidKeyException("Test failure")); + mockedCipher.when(() -> Cipher.getInstance(TEST_WRAP_ALGORITHM)) + .thenReturn(mockCipher); + + try { + AndroidKeyStoreUtil.unwrap( + TEST_WRAPPED_KEY_BYTES, + TEST_KEY_ALGORITHM, + mKeyPair, + TEST_WRAP_ALGORITHM + ); + fail("Should have thrown ClientException"); + } catch (ClientException e) { + assertEquals("INVALID_KEY", e.getErrorCode()); + verify(mockCipher, times(3)).init(Cipher.UNWRAP_MODE, mKeyPair.getPrivate()); + } + } + } + + @Test + public void testUnwrap_NoSuchAlgorithm() throws Exception { + try (MockedStatic 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()); + } + } + } +} diff --git a/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoader.java b/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoader.java index 6e4495fdbd..2f151db8e6 100644 --- a/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoader.java +++ b/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoader.java @@ -54,6 +54,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; @@ -98,6 +100,10 @@ public class AndroidWrappedKeyLoader extends AES256KeyLoader { // Exposed for testing only. /* package */ static final int KEY_FILE_SIZE = 1024; + /** + * SecretKey cache. Maps wrapped secret key file path to the cached data containing the SecretKey. + */ + private static final ConcurrentMap> sKeyCacheMap = new ConcurrentHashMap<>(); private final Context mContext; /** @@ -111,21 +117,15 @@ public class AndroidWrappedKeyLoader extends AES256KeyLoader { */ private final String mFilePath; - private final CachedData mKeyCache = new CachedData() { - @Override - public SecretKey getData() { - if (!sSkipKeyInvalidationCheck && - (!AndroidKeyStoreUtil.canLoadKey(mAlias) || !getKeyFile().exists())) { - this.clear(); - } - return super.getData(); - } - }; - // Exposed for testing only. - @NonNull + @Nullable /* package */ CachedData getKeyCache() { - return mKeyCache; + return sKeyCacheMap.get(mFilePath); + } + + // Exposed for testing only. + /* package */ void clearKeyCache() { + sKeyCacheMap.remove(mFilePath); } /** @@ -162,18 +162,28 @@ 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(); + final CachedData keyCache = sKeyCacheMap.computeIfAbsent(mFilePath, k -> new CachedData() {}); + + if (keyCache.getData() != null) { + return keyCache.getData(); } + Logger.info(methodTag, "Key cache is empty, loading key from storage"); + SecretKey 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); + keyCache.setData(key); return key; } @@ -213,7 +223,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; } @@ -417,7 +427,7 @@ private void recordKeyGenerationTime(long keypairGenStartTime) { public void deleteSecretKeyFromStorage() throws ClientException { AndroidKeyStoreUtil.deleteKey(mAlias); FileUtil.deleteFile(getKeyFile()); - mKeyCache.clear(); + sKeyCacheMap.remove(mFilePath); } /** diff --git a/common/src/main/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtil.java b/common/src/main/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtil.java index c4092efb33..1d68d290f7 100644 --- a/common/src/main/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtil.java +++ b/common/src/main/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtil.java @@ -87,20 +87,19 @@ 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 final int MAX_UNWRAP_RETRIES = 3; + private static final long INITIAL_BACKOFF_MS = 100; + 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; } /** @@ -210,7 +209,7 @@ public static synchronized boolean canLoadKey(@NonNull final String keyAlias) { @Nullable public static synchronized KeyPair readKey(@NonNull final String keyAlias) throws ClientException { - final String methodTag = TAG + ":readKeyPair"; + final String methodTag = TAG + ":readKey"; final Throwable exception; final String errCode; @@ -307,8 +306,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) { @@ -407,7 +405,7 @@ public static synchronized byte[] wrap(@NonNull final SecretKey key, * @param wrapAlgorithm the algorithm used to wrap the key. * @return the unwrapped key. */ - public static synchronized SecretKey unwrap(@NonNull final byte[] wrappedKeyBlob, + public static synchronized SecretKey unwrap(final byte[] wrappedKeyBlob, @NonNull final String wrappedKeyAlgorithm, @NonNull final KeyPair keyPairForUnwrapping, @NonNull final String wrapAlgorithm) throws ClientException { @@ -415,9 +413,7 @@ public static synchronized SecretKey unwrap(@NonNull final byte[] wrappedKeyBlob final Throwable exception; final String errCode; try { - final Cipher wrapCipher = Cipher.getInstance(wrapAlgorithm); - wrapCipher.init(Cipher.UNWRAP_MODE, keyPairForUnwrapping.getPrivate()); - return (SecretKey) wrapCipher.unwrap(wrappedKeyBlob, wrappedKeyAlgorithm, Cipher.SECRET_KEY); + return unwrapWithRetry(wrappedKeyBlob, wrappedKeyAlgorithm, keyPairForUnwrapping, wrapAlgorithm); } catch (final IllegalArgumentException e) { // There is issue with Android KeyStore when lock screen type is changed which could // potentially wipe out keystore. @@ -450,15 +446,6 @@ public static synchronized SecretKey unwrap(@NonNull final byte[] wrappedKeyBlob exception.getMessage(), exception ); - if (exception instanceof InvalidKeyException) { - final Attributes attributes = Attributes.builder() - .put(AttributeName.keystore_operation.name(), "unwrap") - .put(AttributeName.error_code.name(), errCode) - .put(AttributeName.error_type.name(), clientException.getClass().getSimpleName()) - .put(AttributeName.keystore_exception_stack_trace.name(), ThrowableUtil.getStackTraceAsString(clientException)) - .build(); - sFailedAndroidKeyStoreUnwrapOperationCount.add(1, attributes); - } Logger.error( methodTag, errCode, @@ -468,4 +455,49 @@ public static synchronized SecretKey unwrap(@NonNull final byte[] wrappedKeyBlob throw clientException; } + /** + * Internal method to handle unwrap retries with exponential backoff + */ + @NonNull + private static SecretKey unwrapWithRetry(final byte[] wrappedKeyBlob, + @NonNull final String wrappedKeyAlgorithm, + @NonNull final KeyPair keyPairForUnwrapping, + @NonNull final String wrapAlgorithm) throws NoSuchPaddingException, NoSuchAlgorithmException, InvalidKeyException { + final String methodTag = TAG + ":unwrapWithRetry"; + int attempt = 0; + do { + attempt++; + try { + final Cipher wrapCipher = Cipher.getInstance(wrapAlgorithm); + wrapCipher.init(Cipher.UNWRAP_MODE, keyPairForUnwrapping.getPrivate()); + return (SecretKey) wrapCipher.unwrap(wrappedKeyBlob, wrappedKeyAlgorithm, Cipher.SECRET_KEY); + } catch (final InvalidKeyException e) { + final Attributes attributes = Attributes.builder() + .put(AttributeName.keystore_operation.name(), "unwrap") + .put(AttributeName.error_code.name(), INVALID_KEY) + .put(AttributeName.error_type.name(), e.getClass().getSimpleName()) + .put(AttributeName.keystore_exception_stack_trace.name(), ThrowableUtil.getStackTraceAsString(e)) + .build(); + sFailedAndroidKeyStoreUnwrapOperationCount.add(1, attributes); + + long backoffMs = INITIAL_BACKOFF_MS * (1L << attempt); + Logger.error(methodTag, + String.format("InvalidKeyException encountered, (attempt %d/%d) failed. Retrying in %d ms ", attempt, MAX_UNWRAP_RETRIES, backoffMs), + e); + if (attempt == MAX_UNWRAP_RETRIES) { + throw e; + } else { + try { + Thread.sleep(backoffMs); + } catch (final InterruptedException ie) { + Thread.currentThread().interrupt(); + Logger.warn(methodTag, "Thread was interrupted during backoff, rethrowing"); + throw e; + } + } + } + } while (attempt < MAX_UNWRAP_RETRIES); + throw new IllegalStateException("This code should not be reachable"); + } + } diff --git a/common4j/src/main/com/microsoft/identity/common/java/crypto/StorageEncryptionManager.java b/common4j/src/main/com/microsoft/identity/common/java/crypto/StorageEncryptionManager.java index cd7fcc5dc2..035c264174 100644 --- a/common4j/src/main/com/microsoft/identity/common/java/crypto/StorageEncryptionManager.java +++ b/common4j/src/main/com/microsoft/identity/common/java/crypto/StorageEncryptionManager.java @@ -223,8 +223,7 @@ public byte[] decrypt(final byte[] cipherText) throws ClientException { try { return decryptWithSecretKey(dataBytes, keyLoader); } catch (final Throwable e) { - Logger.warn(methodTag, "Failed to decrypt with key:" + keyLoader.getAlias() + - " thumbprint : " + KeyUtil.getKeyThumbPrint(keyLoader)); + Logger.warn(methodTag, "Failed to decrypt with key:" + keyLoader.getAlias()); suppressedException.add(e); } } From c88538f940064f2a51282d262e3db2d0f15acf27 Mon Sep 17 00:00:00 2001 From: Mohit Date: Thu, 29 May 2025 22:34:21 -0700 Subject: [PATCH 2/9] fix lint --- changelog.txt | 2 +- .../crypto/AndroidWrappedKeyLoaderTest.java | 13 +++++------- .../crypto/AndroidWrappedKeyLoader.java | 20 +++++++++---------- 3 files changed, 15 insertions(+), 20 deletions(-) diff --git a/changelog.txt b/changelog.txt index 62e6576590..88b5022a77 100644 --- a/changelog.txt +++ b/changelog.txt @@ -1,6 +1,6 @@ vNext ---------- -- [PATCH] Fix caching of secret key and add retries for InvalidKeyException during unwrap +- [PATCH] Fix caching of secret key and add retries for InvalidKeyException during unwrap (#2659) - [MINOR] Add Get AAD device ID API for resource account (#2644) - [MINOR] Expose a JavaScript API in brokered Webviews to facilitate Improved Same Device NumberMatch (#2617) - [MINOR] Add API for resource account provisioning (API only) (#2640) diff --git a/common/src/androidTest/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoaderTest.java b/common/src/androidTest/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoaderTest.java index 7d59cbb5d5..6fc5ee42e0 100644 --- a/common/src/androidTest/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoaderTest.java +++ b/common/src/androidTest/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoaderTest.java @@ -31,7 +31,6 @@ import com.microsoft.identity.common.internal.util.AndroidKeyStoreUtil; import com.microsoft.identity.common.java.crypto.key.AES256KeyLoader; import com.microsoft.identity.common.java.exception.ClientException; -import com.microsoft.identity.common.java.util.CachedData; import com.microsoft.identity.common.java.util.FileUtil; import org.junit.Assert; @@ -151,9 +150,7 @@ public void testLoadKey() throws ClientException { final AndroidWrappedKeyLoader keyLoader = new AndroidWrappedKeyLoader(MOCK_KEY_ALIAS, MOCK_KEY_FILE_PATH, context); final SecretKey secretKey = keyLoader.getKey(); - final CachedData cachedKey = keyLoader.getKeyCache(); - Assert.assertNotNull(cachedKey); - final SecretKey key = cachedKey.getData(); + final SecretKey key = keyLoader.getKeyFromCache(); Assert.assertNotNull(key); Assert.assertEquals(AES256KeyLoader.AES_ALGORITHM, secretKey.getAlgorithm()); Assert.assertArrayEquals(secretKey.getEncoded(), key.getEncoded()); @@ -237,7 +234,7 @@ public void testPerf_NoCachedKey() throws ClientException { long timeStartLoopNotCached = System.nanoTime(); for (int i = 0; i < 100; i++) { - keyLoader.clearKeyCache(); + keyLoader.clearKeyFromCache(); keyLoader.getKey(); } long timeFinishLoopNotCached = System.nanoTime(); @@ -255,7 +252,7 @@ public void testLoadDeletedKeyStoreKey() throws ClientException { AndroidKeyStoreUtil.deleteKey(MOCK_KEY_ALIAS); // Cached key also be wiped. - final CachedData key = keyLoader.getKeyCache(); + final SecretKey key = keyLoader.getKeyFromCache(); Assert.assertNull(key); } @@ -266,7 +263,7 @@ public void testLoadDeletedKeyFile() throws ClientException { FileUtil.deleteFile(getKeyFile()); // Cached key also be wiped. - final CachedData key = keyLoader.getKeyCache(); + final SecretKey key = keyLoader.getKeyFromCache(); Assert.assertNull(key); } @@ -365,7 +362,7 @@ private AndroidWrappedKeyLoader initKeyLoaderWithKeyEntry() throws ClientExcepti final AndroidWrappedKeyLoader keyLoader = new AndroidWrappedKeyLoader(MOCK_KEY_ALIAS, MOCK_KEY_FILE_PATH, context); final SecretKey key = keyLoader.getKey(); Assert.assertNotNull(key); - Assert.assertNotNull(keyLoader.getKeyCache()); + Assert.assertNotNull(keyLoader.getKeyFromCache()); return keyLoader; } } diff --git a/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoader.java b/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoader.java index 2f151db8e6..986265ecd9 100644 --- a/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoader.java +++ b/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoader.java @@ -41,7 +41,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; @@ -101,9 +100,9 @@ public class AndroidWrappedKeyLoader extends AES256KeyLoader { /* package */ static final int KEY_FILE_SIZE = 1024; /** - * SecretKey cache. Maps wrapped secret key file path to the cached data containing the SecretKey. + * SecretKey cache. Maps wrapped secret key file path to the SecretKey. */ - private static final ConcurrentMap> sKeyCacheMap = new ConcurrentHashMap<>(); + private static final ConcurrentMap sKeyCacheMap = new ConcurrentHashMap<>(); private final Context mContext; /** @@ -119,12 +118,12 @@ public class AndroidWrappedKeyLoader extends AES256KeyLoader { // Exposed for testing only. @Nullable - /* package */ CachedData getKeyCache() { + /* package */ SecretKey getKeyFromCache() { return sKeyCacheMap.get(mFilePath); } // Exposed for testing only. - /* package */ void clearKeyCache() { + /* package */ void clearKeyFromCache() { sKeyCacheMap.remove(mFilePath); } @@ -168,13 +167,12 @@ public synchronized SecretKey getKey() throws ClientException { sKeyCacheMap.remove(mFilePath); } - final CachedData keyCache = sKeyCacheMap.computeIfAbsent(mFilePath, k -> new CachedData() {}); - - if (keyCache.getData() != null) { - return keyCache.getData(); + SecretKey cachedKey = sKeyCacheMap.get(mFilePath); + if (cachedKey != null) { + return cachedKey; } - Logger.info(methodTag, "Key cache is empty, loading key from storage"); + Logger.info(methodTag, "Key not in cache is empty, loading key from storage"); SecretKey key = readSecretKeyFromStorage(); // If key doesn't exist, generate a new one. @@ -183,7 +181,7 @@ public synchronized SecretKey getKey() throws ClientException { key = generateRandomKey(); } - keyCache.setData(key); + sKeyCacheMap.put(mFilePath, key); return key; } From 674f6db364335e54a5d84fbd505061537d4b721c Mon Sep 17 00:00:00 2001 From: Mohit Date: Fri, 30 May 2025 11:00:09 -0700 Subject: [PATCH 3/9] update local var --- .../identity/common/crypto/AndroidWrappedKeyLoader.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoader.java b/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoader.java index 986265ecd9..46eb8cfeb8 100644 --- a/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoader.java +++ b/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyLoader.java @@ -167,13 +167,13 @@ public synchronized SecretKey getKey() throws ClientException { sKeyCacheMap.remove(mFilePath); } - SecretKey cachedKey = sKeyCacheMap.get(mFilePath); - if (cachedKey != null) { - return cachedKey; + SecretKey key = sKeyCacheMap.get(mFilePath); + if (key != null) { + return key; } Logger.info(methodTag, "Key not in cache is empty, loading key from storage"); - SecretKey key = readSecretKeyFromStorage(); + key = readSecretKeyFromStorage(); // If key doesn't exist, generate a new one. if (key == null) { From 865060447de1cbe4359b17814138854d42ba71f3 Mon Sep 17 00:00:00 2001 From: Mohit Date: Tue, 8 Jul 2025 15:26:32 -0700 Subject: [PATCH 4/9] fix merge --- .../crypto/AndroidWrappedKeyProviderTest.java | 14 ++--- .../util/AndroidKeyStoreUtilTest.java | 60 +------------------ .../internal/util/AndroidKeyStoreUtil.java | 13 ++-- 3 files changed, 12 insertions(+), 75 deletions(-) diff --git a/common/src/androidTest/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProviderTest.java b/common/src/androidTest/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProviderTest.java index e23cd4adb2..c3bfa44ea5 100644 --- a/common/src/androidTest/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProviderTest.java +++ b/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,7 +257,7 @@ 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); } @@ -265,7 +265,7 @@ private AndroidWrappedKeyProvider initkeyProviderWithKeyEntry() throws ClientExc 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; } } diff --git a/common/src/androidTest/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtilTest.java b/common/src/androidTest/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtilTest.java index 02ed1fcb6c..c23139c27a 100644 --- a/common/src/androidTest/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtilTest.java +++ b/common/src/androidTest/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtilTest.java @@ -28,12 +28,9 @@ import static org.junit.Assert.fail; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; - -import android.security.keystore.KeyGenParameterSpec; -import android.security.keystore.KeyProperties; +import static org.mockito.Mockito.when; import androidx.test.ext.junit.runners.AndroidJUnit4; @@ -45,13 +42,9 @@ import org.mockito.MockedStatic; import org.mockito.Mockito; -import java.security.InvalidKeyException; import java.security.KeyPair; import java.security.KeyPairGenerator; -import java.security.KeyStore; import java.security.NoSuchAlgorithmException; -import java.security.PrivateKey; -import java.security.PublicKey; import javax.crypto.Cipher; import javax.crypto.SecretKey; @@ -102,57 +95,6 @@ public void testUnwrap_Success() throws Exception { } } - @Test - public void testUnwrap_WithRetries() throws Exception { - // Mock Cipher to fail twice with InvalidKeyException then succeed - try (MockedStatic mockedCipher = Mockito.mockStatic(Cipher.class)) { - Cipher mockCipher = mock(Cipher.class); - when(mockCipher.unwrap(any(byte[].class), any(String.class), any(int.class))) - .thenThrow(new InvalidKeyException("Test failure 1")) - .thenThrow(new InvalidKeyException("Test failure 2")) - .thenReturn(mSecretKey); - mockedCipher.when(() -> Cipher.getInstance(TEST_WRAP_ALGORITHM)) - .thenReturn(mockCipher); - - // Test unwrap with retries - SecretKey result = AndroidKeyStoreUtil.unwrap( - TEST_WRAPPED_KEY_BYTES, - TEST_KEY_ALGORITHM, - mKeyPair, - TEST_WRAP_ALGORITHM - ); - - assertNotNull(result); - assertEquals(mSecretKey, result); - verify(mockCipher, times(3)).init(Cipher.UNWRAP_MODE, mKeyPair.getPrivate()); - } - } - - @Test - public void testUnwrap_MaxRetriesExceeded() throws Exception { - // Mock Cipher to always fail with InvalidKeyException - try (MockedStatic mockedCipher = Mockito.mockStatic(Cipher.class)) { - Cipher mockCipher = mock(Cipher.class); - when(mockCipher.unwrap(any(byte[].class), any(String.class), any(int.class))) - .thenThrow(new InvalidKeyException("Test failure")); - mockedCipher.when(() -> Cipher.getInstance(TEST_WRAP_ALGORITHM)) - .thenReturn(mockCipher); - - try { - AndroidKeyStoreUtil.unwrap( - TEST_WRAPPED_KEY_BYTES, - TEST_KEY_ALGORITHM, - mKeyPair, - TEST_WRAP_ALGORITHM - ); - fail("Should have thrown ClientException"); - } catch (ClientException e) { - assertEquals("INVALID_KEY", e.getErrorCode()); - verify(mockCipher, times(3)).init(Cipher.UNWRAP_MODE, mKeyPair.getPrivate()); - } - } - } - @Test public void testUnwrap_NoSuchAlgorithm() throws Exception { try (MockedStatic mockedCipher = Mockito.mockStatic(Cipher.class)) { diff --git a/common/src/main/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtil.java b/common/src/main/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtil.java index c10544c6be..84abd4f7d1 100644 --- a/common/src/main/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtil.java +++ b/common/src/main/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtil.java @@ -87,8 +87,6 @@ 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" @@ -96,11 +94,9 @@ private AndroidKeyStoreUtil() { 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) { From 14828d3f314b6f176959d99950019e37b4097778 Mon Sep 17 00:00:00 2001 From: Mohit Date: Wed, 9 Jul 2025 14:33:35 -0700 Subject: [PATCH 5/9] Update common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java Co-authored-by: Sowmya Malayanur <69237821+somalaya@users.noreply.github.com> --- .../identity/common/crypto/AndroidWrappedKeyProvider.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java b/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java index 78c737cfdf..addec4abc3 100644 --- a/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java +++ b/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java @@ -183,7 +183,7 @@ public synchronized SecretKey getKey() throws ClientException { return key; } - Logger.info(methodTag, "Key not in cache is empty, loading key from storage"); + Logger.info(methodTag, "Key not in cache or cache is empty, loading key from storage"); key = readSecretKeyFromStorage(); // If key doesn't exist, generate a new one. From e57a68496d8504456d08f0dc9bdca9ebae5b3a01 Mon Sep 17 00:00:00 2001 From: Mohit Date: Wed, 9 Jul 2025 14:35:00 -0700 Subject: [PATCH 6/9] addressed comment --- .../identity/common/crypto/AndroidWrappedKeyProvider.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java b/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java index addec4abc3..bff8240ae6 100644 --- a/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java +++ b/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java @@ -31,6 +31,7 @@ import android.security.keystore.KeyProperties; import androidx.annotation.RequiresApi; +import androidx.annotation.VisibleForTesting; import com.microsoft.identity.common.internal.util.AndroidKeyStoreUtil; import com.microsoft.identity.common.java.controllers.ExceptionAdapter; @@ -129,6 +130,7 @@ public class AndroidWrappedKeyProvider implements ISecretKeyProvider { // Exposed for testing only. @Nullable + @VisibleForTesting /* package */ SecretKey getKeyFromCache() { return sKeyCacheMap.get(mFilePath); } From 49939df2371be4dadd95dc4b98465ef2782d1951 Mon Sep 17 00:00:00 2001 From: Mohit Date: Wed, 9 Jul 2025 14:35:46 -0700 Subject: [PATCH 7/9] addressed comment --- .../identity/common/crypto/AndroidWrappedKeyProvider.java | 1 + 1 file changed, 1 insertion(+) diff --git a/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java b/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java index bff8240ae6..1cd50a8f85 100644 --- a/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java +++ b/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java @@ -136,6 +136,7 @@ public class AndroidWrappedKeyProvider implements ISecretKeyProvider { } // Exposed for testing only. + @VisibleForTesting /* package */ void clearKeyFromCache() { sKeyCacheMap.remove(mFilePath); } From 49dcf718050269e575002335a5a009c73ce12755 Mon Sep 17 00:00:00 2001 From: Mohit Date: Mon, 14 Jul 2025 10:52:05 -0700 Subject: [PATCH 8/9] addressed comment --- .../common/crypto/AndroidWrappedKeyProvider.java | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java b/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java index 1cd50a8f85..088f728fac 100644 --- a/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java +++ b/common/src/main/java/com/microsoft/identity/common/crypto/AndroidWrappedKeyProvider.java @@ -132,6 +132,12 @@ public class AndroidWrappedKeyProvider implements ISecretKeyProvider { @Nullable @VisibleForTesting /* package */ SecretKey getKeyFromCache() { + final String methodTag = TAG + ":getKeyFromCache"; + if (!sSkipKeyInvalidationCheck && + (!AndroidKeyStoreUtil.canLoadKey(mAlias) || !this.getKeyFile().exists())) { + Logger.warn(methodTag, "Key is invalid, removing from cache"); + clearKeyFromCache(); + } return sKeyCacheMap.get(mFilePath); } @@ -176,12 +182,8 @@ public String getKeyTypeIdentifier() { @NonNull public synchronized SecretKey getKey() throws ClientException { final String methodTag = TAG + ":getKey"; - if (!sSkipKeyInvalidationCheck && - (!AndroidKeyStoreUtil.canLoadKey(mAlias) || !this.getKeyFile().exists())) { - sKeyCacheMap.remove(mFilePath); - } - SecretKey key = sKeyCacheMap.get(mFilePath); + SecretKey key = this.getKeyFromCache(); if (key != null) { return key; } From c2bb1d11f5928a8c0d218aa84550da97aea3c432 Mon Sep 17 00:00:00 2001 From: Mohit Chandwani Date: Mon, 14 Jul 2025 11:59:56 -0700 Subject: [PATCH 9/9] Fix UT --- .../util/AndroidKeyStoreUtilTest.java | 88 ++++++++++--------- 1 file changed, 47 insertions(+), 41 deletions(-) diff --git a/common/src/androidTest/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtilTest.java b/common/src/androidTest/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtilTest.java index c23139c27a..18c0e125a3 100644 --- a/common/src/androidTest/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtilTest.java +++ b/common/src/androidTest/java/com/microsoft/identity/common/internal/util/AndroidKeyStoreUtilTest.java @@ -23,14 +23,11 @@ package com.microsoft.identity.common.internal.util; +import static com.microsoft.identity.common.java.exception.ClientException.INVALID_KEY; +import static com.microsoft.identity.common.java.exception.ClientException.NO_SUCH_ALGORITHM; 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; @@ -39,14 +36,10 @@ 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; @@ -72,46 +65,59 @@ public void setUp() throws Exception { } @Test - public void testUnwrap_Success() throws Exception { - // Mock Cipher for successful unwrap - try (MockedStatic 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, + public void testUnwrap_Success() throws ClientException { + // Test successful wrapping + final byte[] wrapped = AndroidKeyStoreUtil.wrap( + mSecretKey, + mKeyPair, + TEST_WRAP_ALGORITHM + ); + + // Test successful unwrap + final SecretKey unwrapped = AndroidKeyStoreUtil.unwrap( + wrapped, + TEST_KEY_ALGORITHM, + mKeyPair, + TEST_WRAP_ALGORITHM + ); + + assertNotNull(unwrapped); + assertEquals(mSecretKey, unwrapped); + } + + @Test + public void testUnwrap_NoSuchAlgorithm() { + try { + final byte[] wrapped = AndroidKeyStoreUtil.wrap( + mSecretKey, mKeyPair, TEST_WRAP_ALGORITHM ); - assertNotNull(result); - assertEquals(mSecretKey, result); - verify(mockCipher, times(1)).init(Cipher.UNWRAP_MODE, mKeyPair.getPrivate()); + final SecretKey secretKey = AndroidKeyStoreUtil.unwrap( + wrapped, + TEST_KEY_ALGORITHM, + mKeyPair, + "NoAlg" + ); + fail("Should have thrown ClientException"); + } catch (final ClientException e) { + assertEquals(NO_SUCH_ALGORITHM, e.getErrorCode()); } } @Test - public void testUnwrap_NoSuchAlgorithm() throws Exception { - try (MockedStatic 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()); - } + public void testUnwrap_InvalidKey() { + try { + AndroidKeyStoreUtil.unwrap( + TEST_WRAPPED_KEY_BYTES, + TEST_KEY_ALGORITHM, + mKeyPair, + TEST_WRAP_ALGORITHM + ); + fail("Should have thrown ClientException"); + } catch (final ClientException e) { + assertEquals(INVALID_KEY, e.getErrorCode()); } } }