Fixing bug calling get credentials multiple times after migration

Author: utkrishtsahu

auth0/src/main/java/com/auth0/android/authentication/storage/CryptoUtil.java

@@ -429,8 +429,8 @@ byte[] RSAEncrypt(byte[] decryptedInput) throws IncompatibleDeviceException, Cry
 
     /**
      * Attempts to migrate legacy PKCS1-encrypted AES key to OAEP format.
-     * This method tries to decrypt the AES key using legacy PKCS1 padding.
-     * If successful, it deletes the old keys so fresh OAEP keys can be generated.
+     * This method tries to decrypt the AES key using legacy PKCS1 padding,
+     * then re-encrypts it with OAEP and stores it for future use.
      *
      * @param encryptedAESBytes the encrypted AES key bytes
      * @return the decrypted AES key if migration succeeds, or null if migration fails
@@ -455,9 +455,14 @@ private byte[] attemptPKCS1Migration(byte[] encryptedAESBytes) {
             }
 
             Log.d(TAG, "PKCS1 migration successful - deleting old keys");
+            
             deleteRSAKeys();
-            deleteAESKeys();
 
+            byte[] encryptedAESWithOAEP = RSAEncrypt(decryptedAESKey);
+            String encodedEncryptedAES = new String(Base64.encode(encryptedAESWithOAEP, Base64.DEFAULT), StandardCharsets.UTF_8);
+            storage.store(KEY_ALIAS, encodedEncryptedAES);
+            
+            Log.d(TAG, "AES key re-encrypted with OAEP and stored");
             return decryptedAESKey;
 
         } catch (BadPaddingException | IllegalBlockSizeException e) {
@@ -466,6 +471,8 @@ private byte[] attemptPKCS1Migration(byte[] encryptedAESBytes) {
                  NoSuchAlgorithmException | UnrecoverableEntryException |
                  NoSuchPaddingException | InvalidKeyException e) {
             Log.e(TAG, "Migration failed due to key access error.", e);
+        } catch (CryptoException e) {
+            Log.e(TAG, "Failed to re-encrypt AES key with OAEP.", e);
         }
         return null;
     }
@@ -599,15 +606,15 @@ private byte[] tryMigrateLegacyAESKey() {
         try {
             byte[] encryptedOldAESBytes = Base64.decode(encodedOldAES, Base64.DEFAULT);
             KeyStore.PrivateKeyEntry rsaKeyEntry = getRSAKeyEntry();
-            
+
             byte[] decryptedAESKey = RSADecryptLegacyPKCS1(encryptedOldAESBytes, rsaKeyEntry.getPrivateKey());
 
             // Re-encrypt with OAEP and store at new location
             byte[] encryptedAESWithOAEP = RSAEncrypt(decryptedAESKey);
             String newEncodedEncryptedAES = new String(Base64.encode(encryptedAESWithOAEP, Base64.DEFAULT), StandardCharsets.UTF_8);
             storage.store(KEY_ALIAS, newEncodedEncryptedAES);
             storage.remove(OLD_KEY_ALIAS);
-            
+
             Log.d(TAG, "Legacy AES key migrated successfully");
             return decryptedAESKey;
         } catch (Exception e) {

auth0/src/test/java/com/auth0/android/authentication/storage/CryptoUtilTest.java

@@ -1715,12 +1715,16 @@ public void shouldDetectAndMigratePKCS1KeyToOAEP() throws Exception {
 
         byte[] aesKeyBytes = new byte[]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32};
         byte[] encryptedAESKeyPKCS1 = new byte[]{20, 21, 22, 23, 24};
+        byte[] encryptedAESKeyOAEP = new byte[]{30, 31, 32, 33, 34};
         String encodedEncryptedAESPKCS1 = "pkcs1_encrypted_key";
+        String encodedEncryptedAESOAEP = "oaep_encrypted_key";
 
         when(storage.retrieveString(eq(KEY_ALIAS))).thenReturn(encodedEncryptedAESPKCS1);
         when(storage.retrieveString(eq(OLD_KEY_ALIAS))).thenReturn(null);
         PowerMockito.mockStatic(Base64.class);
         PowerMockito.when(Base64.decode(encodedEncryptedAESPKCS1, Base64.DEFAULT)).thenReturn(encryptedAESKeyPKCS1);
+        PowerMockito.when(Base64.encode(encryptedAESKeyOAEP, Base64.DEFAULT))
+            .thenReturn(encodedEncryptedAESOAEP.getBytes(StandardCharsets.UTF_8));
 
         IncompatibleDeviceException incompatibleException = new IncompatibleDeviceException(
             new KeyStoreException("Incompatible padding mode")
@@ -1730,24 +1734,27 @@ public void shouldDetectAndMigratePKCS1KeyToOAEP() throws Exception {
         when(keyStore.containsAlias(KEY_ALIAS)).thenReturn(true);
         KeyStore.PrivateKeyEntry mockKeyEntry = mock(KeyStore.PrivateKeyEntry.class);
         PrivateKey mockPrivateKey = mock(PrivateKey.class);
+        Certificate mockCertificate = mock(Certificate.class);
+        PublicKey mockPublicKey = mock(PublicKey.class);
         when(mockKeyEntry.getPrivateKey()).thenReturn(mockPrivateKey);
+        when(mockKeyEntry.getCertificate()).thenReturn(mockCertificate);
+        when(mockCertificate.getPublicKey()).thenReturn(mockPublicKey);
         when(keyStore.getEntry(eq(KEY_ALIAS), nullable(KeyStore.ProtectionParameter.class)))
             .thenReturn(mockKeyEntry);
 
         when(rsaPkcs1Cipher.doFinal(encryptedAESKeyPKCS1)).thenReturn(aesKeyBytes);
+        
+        doReturn(encryptedAESKeyOAEP).when(cryptoUtil).RSAEncrypt(aesKeyBytes);
 
         byte[] result = cryptoUtil.getAESKey();
 
         assertThat(result, is(aesKeyBytes));
 
-        // Verify PKCS1 cipher was used to decrypt the legacy key
         Mockito.verify(rsaPkcs1Cipher).init(Cipher.DECRYPT_MODE, mockPrivateKey);
         Mockito.verify(rsaPkcs1Cipher).doFinal(encryptedAESKeyPKCS1);
 
-        // Simplified migration: RSA and AES keys are deleted (not re-encrypted)
-        // Next operation will generate fresh OAEP keys
         Mockito.verify(keyStore).deleteEntry(KEY_ALIAS);
-        Mockito.verify(storage).remove(KEY_ALIAS);
+        Mockito.verify(storage).store(KEY_ALIAS, encodedEncryptedAESOAEP);
     }
 
     @Test
@@ -1861,13 +1868,17 @@ public void shouldAllowMultipleRetrievalsAfterMigration() throws Exception {
 
         byte[] aesKeyBytes = new byte[]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32};
         byte[] encryptedAESKeyPKCS1 = new byte[]{20, 21, 22, 23, 24};
+        byte[] encryptedAESKeyOAEP = new byte[]{30, 31, 32, 33, 34};
         String encodedEncryptedAESPKCS1 = "pkcs1_encrypted_key";
+        String encodedEncryptedAESOAEP = "oaep_encrypted_key";
 
         // First retrieval - migration happens, returns decrypted key
         when(storage.retrieveString(eq(KEY_ALIAS))).thenReturn(encodedEncryptedAESPKCS1);
         when(storage.retrieveString(eq(OLD_KEY_ALIAS))).thenReturn(null);
         PowerMockito.mockStatic(Base64.class);
         PowerMockito.when(Base64.decode(encodedEncryptedAESPKCS1, Base64.DEFAULT)).thenReturn(encryptedAESKeyPKCS1);
+        PowerMockito.when(Base64.encode(encryptedAESKeyOAEP, Base64.DEFAULT))
+            .thenReturn(encodedEncryptedAESOAEP.getBytes(StandardCharsets.UTF_8));
 
         IncompatibleDeviceException incompatibleException = new IncompatibleDeviceException(
             new KeyStoreException("Incompatible padding mode")
@@ -1877,18 +1888,25 @@ public void shouldAllowMultipleRetrievalsAfterMigration() throws Exception {
         when(keyStore.containsAlias(KEY_ALIAS)).thenReturn(true);
         KeyStore.PrivateKeyEntry mockKeyEntry = mock(KeyStore.PrivateKeyEntry.class);
         PrivateKey mockPrivateKey = mock(PrivateKey.class);
+        Certificate mockCertificate = mock(Certificate.class);
+        PublicKey mockPublicKey = mock(PublicKey.class);
         when(mockKeyEntry.getPrivateKey()).thenReturn(mockPrivateKey);
+        when(mockKeyEntry.getCertificate()).thenReturn(mockCertificate);
+        when(mockCertificate.getPublicKey()).thenReturn(mockPublicKey);
         when(keyStore.getEntry(eq(KEY_ALIAS), nullable(KeyStore.ProtectionParameter.class)))
             .thenReturn(mockKeyEntry);
 
         when(rsaPkcs1Cipher.doFinal(encryptedAESKeyPKCS1)).thenReturn(aesKeyBytes);
+        
+        // Mock RSAEncrypt for re-encrypting with OAEP after migration
+        doReturn(encryptedAESKeyOAEP).when(cryptoUtil).RSAEncrypt(aesKeyBytes);
 
         byte[] result1 = cryptoUtil.getAESKey();
         assertThat(result1, is(aesKeyBytes));
 
-        // Simplified migration: keys are deleted, not re-stored
+        // Migration should delete old keys and store re-encrypted AES key
         Mockito.verify(keyStore).deleteEntry(KEY_ALIAS);
-        Mockito.verify(storage).remove(KEY_ALIAS);
+        Mockito.verify(storage).store(KEY_ALIAS, encodedEncryptedAESOAEP);
     }
 
     @Test