Batch fetching of key descriptors from Keystore

Change interaction with Keystore2 in the following manner:
* Return an enumerator over the entries in Keystore2 rather than
  attempting to get all of them into one single data structure.
* Use a new Keystore2 method for getting the count of entries
  rather than count the size of the array returned.

The enumerator reads a batch of key descriptors from Keystore2.
Once the batch has been exhausted, the enumerator added asks
Keystore2 for the next batch of keys starting with the last
alias it has processed, until it receives an empty array.

Bug: 222287335
Test: atest KeystoreTests
Change-Id: I309b3188df998825557a3c5e6d777b1c0807a924

3 files changed, 241 insertions, 18 deletions

diff --git a/keystore/java/android/security/KeyStore2.java b/keystore/java/android/security/KeyStore2.java
index c2cd6ffc622c..74597c5cd874 100644
--- a/keystore/java/android/security/KeyStore2.java
+++ b/keystore/java/android/security/KeyStore2.java
@@ -161,6 +161,15 @@ public class KeyStore2 {
     }
 
     /**
+     * List all entries in the keystore for in the given namespace.
+     */
+    public KeyDescriptor[] listBatch(int domain, long namespace, String startPastAlias)
+            throws KeyStoreException {
+        return handleRemoteExceptionWithRetry(
+                (service) -> service.listEntriesBatched(domain, namespace, startPastAlias));
+    }
+
+    /**
      * Grant string prefix as used by the keystore boringssl engine. Must be kept in sync
      * with system/security/keystore-engine. Note: The prefix here includes the 0x which
      * std::stringstream used in keystore-engine needs to identify the number as hex represented.
@@ -301,6 +310,13 @@ public class KeyStore2 {
         });
     }
 
+    /**
+     * Returns the number of Keystore entries for a given domain and namespace.
+     */
+    public int getNumberOfEntries(int domain, long namespace) throws KeyStoreException {
+        return handleRemoteExceptionWithRetry((service)
+                -> service.getNumberOfEntries(domain, namespace));
+    }
     protected static void interruptedPreservingSleep(long millis) {
         boolean wasInterrupted = false;
         Calendar calendar = Calendar.getInstance();
diff --git a/keystore/java/android/security/keystore2/AndroidKeyStoreSpi.java b/keystore/java/android/security/keystore2/AndroidKeyStoreSpi.java
index 91f216f1320a..045e318ff513 100644
--- a/keystore/java/android/security/keystore2/AndroidKeyStoreSpi.java
+++ b/keystore/java/android/security/keystore2/AndroidKeyStoreSpi.java
@@ -79,13 +79,11 @@ import java.security.spec.NamedParameterSpec;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.Date;
 import java.util.Enumeration;
-import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
-import java.util.Set;
+import java.util.NoSuchElementException;
 
 import javax.crypto.SecretKey;
 
@@ -1043,26 +1041,22 @@ public class AndroidKeyStoreSpi extends KeyStoreSpi {
         }
     }
 
-    private Set<String> getUniqueAliases() {
+    private KeyDescriptor[] getAliasesBatch(String startPastAlias) {
         try {
-            final KeyDescriptor[] keys = mKeyStore.list(
+            return mKeyStore.listBatch(
                     getTargetDomain(),
-                    mNamespace
+                    mNamespace,
+                    startPastAlias
             );
-            final Set<String> aliases = new HashSet<>(keys.length);
-            for (KeyDescriptor d : keys) {
-                aliases.add(d.alias);
-            }
-            return aliases;
         } catch (android.security.KeyStoreException e) {
             Log.e(TAG, "Failed to list keystore entries.", e);
-            return new HashSet<>();
+            return new KeyDescriptor[0];
         }
     }
 
     @Override
     public Enumeration<String> engineAliases() {
-        return Collections.enumeration(getUniqueAliases());
+        return new KeyEntriesEnumerator();
     }
 
     @Override
@@ -1073,12 +1067,18 @@ public class AndroidKeyStoreSpi extends KeyStoreSpi {
 
         return getKeyMetadata(alias) != null;
     }
-
     @Override
     public int engineSize() {
-        return getUniqueAliases().size();
+        try {
+            return mKeyStore.getNumberOfEntries(
+                    getTargetDomain(),
+                    mNamespace
+            );
+        } catch (android.security.KeyStoreException e) {
+            Log.e(TAG, "Failed to get the number of keystore entries.", e);
+            return 0;
+        }
     }
-
     @Override
     public boolean engineIsKeyEntry(String alias) {
         return isKeyEntry(alias);
@@ -1251,4 +1251,38 @@ public class AndroidKeyStoreSpi extends KeyStoreSpi {
                             + "or TrustedCertificateEntry; was " + entry);
         }
     }
+
+    private class KeyEntriesEnumerator implements Enumeration<String> {
+        private KeyDescriptor[] mCurrentBatch;
+        private int mCurrentEntry = 0;
+        private String mLastAlias = null;
+        private KeyEntriesEnumerator() {
+            getAndValidateNextBatch();
+        }
+
+        private void getAndValidateNextBatch() {
+            mCurrentBatch = getAliasesBatch(mLastAlias);
+            mCurrentEntry = 0;
+        }
+
+        public boolean hasMoreElements() {
+            return (mCurrentBatch != null) && (mCurrentBatch.length > 0);
+        }
+
+        public String nextElement() {
+            if ((mCurrentBatch == null) || (mCurrentBatch.length == 0)) {
+                throw new NoSuchElementException("Error while fetching entries.");
+            }
+            final KeyDescriptor currentEntry = mCurrentBatch[mCurrentEntry];
+            mLastAlias = currentEntry.alias;
+
+            mCurrentEntry++;
+            // This was the last entry in the batch.
+            if (mCurrentEntry >= mCurrentBatch.length) {
+                getAndValidateNextBatch();
+            }
+
+            return mLastAlias;
+        }
+    }
 }
diff --git a/keystore/tests/src/android/security/keystore2/AndroidKeyStoreSpiTest.java b/keystore/tests/src/android/security/keystore2/AndroidKeyStoreSpiTest.java
index f96c39c879fd..1e1f68aaefc1 100644
--- a/keystore/tests/src/android/security/keystore2/AndroidKeyStoreSpiTest.java
+++ b/keystore/tests/src/android/security/keystore2/AndroidKeyStoreSpiTest.java
@@ -17,9 +17,14 @@
 package android.security.keystore2;
 
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertThrows;
 import static org.mockito.Mockito.anyInt;
 import static org.mockito.Mockito.anyLong;
+import static org.mockito.Mockito.anyString;
+import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.isNull;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
 import android.security.KeyStore2;
@@ -36,6 +41,12 @@ import org.junit.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Enumeration;
+import java.util.List;
+import java.util.NoSuchElementException;
+
 public class AndroidKeyStoreSpiTest {
 
     @Mock
@@ -48,14 +59,176 @@ public class AndroidKeyStoreSpiTest {
 
     @Test
     public void testEngineAliasesReturnsEmptySetOnKeyStoreError() throws Exception {
-        when(mKeystore2.list(anyInt(), anyLong()))
+        when(mKeystore2.listBatch(anyInt(), anyLong(), isNull()))
                 .thenThrow(new KeyStoreException(6, "Some Error"));
         AndroidKeyStoreSpi spi = new AndroidKeyStoreSpi();
         spi.initForTesting(mKeystore2);
 
         assertThat("Empty collection expected", !spi.engineAliases().hasMoreElements());
 
-        verify(mKeystore2).list(anyInt(), anyLong());
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), isNull());
+    }
+
+    @Test
+    public void testEngineAliasesCorrectlyListsZeroEntriesEmptyArray() throws Exception {
+        when(mKeystore2.listBatch(anyInt(), anyLong(), anyString()))
+                .thenReturn(new KeyDescriptor[0]);
+        AndroidKeyStoreSpi spi = new AndroidKeyStoreSpi();
+        spi.initForTesting(mKeystore2);
+
+        Enumeration<String> aliases = spi.engineAliases();
+        assertThat("Should not have any elements", !aliases.hasMoreElements());
+        assertThrows("Should have no elements to return", NoSuchElementException.class,
+                () -> aliases.nextElement());
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), isNull());
+    }
+
+    @Test
+    public void testEngineAliasesCorrectlyListsZeroEntriesNullArray() throws Exception {
+        when(mKeystore2.listBatch(anyInt(), anyLong(), anyString()))
+                .thenReturn(null);
+        AndroidKeyStoreSpi spi = new AndroidKeyStoreSpi();
+        spi.initForTesting(mKeystore2);
+
+        Enumeration<String> aliases = spi.engineAliases();
+        assertThat("Should not have any elements", !aliases.hasMoreElements());
+        assertThrows("Should have no elements to return", NoSuchElementException.class,
+                () -> aliases.nextElement());
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), isNull());
+    }
+
+    private static KeyDescriptor newKeyDescriptor(String alias) {
+        KeyDescriptor result = new KeyDescriptor();
+        result.alias = alias;
+        return result;
+    }
+
+    private static KeyDescriptor[] createKeyDescriptorsArray(int numEntries) {
+        KeyDescriptor[] kds = new KeyDescriptor[numEntries];
+        for (int i = 0; i < kds.length; i++) {
+            kds[i] = newKeyDescriptor(String.format("alias-%d", i));
+        }
+
+        return kds;
+    }
+
+    private static void assertAliasListsEqual(
+            KeyDescriptor[] keyDescriptors, Enumeration<String> aliasesEnumerator) {
+        List<String> aliases = Collections.list(aliasesEnumerator);
+        Assert.assertArrayEquals(Arrays.stream(keyDescriptors).map(kd -> kd.alias).toArray(),
+                aliases.toArray());
+    }
+
+    @Test
+    public void testEngineAliasesCorrectlyListsEntriesInASingleBatch() throws Exception {
+        final String alias1 = "testAlias1";
+        final String alias2 = "testAlias2";
+        final String alias3 = "testAlias3";
+        KeyDescriptor[] kds = {newKeyDescriptor(alias1),
+                newKeyDescriptor(alias2), newKeyDescriptor(alias3)};
+        when(mKeystore2.listBatch(anyInt(), anyLong(), eq(null)))
+                .thenReturn(kds);
+        when(mKeystore2.listBatch(anyInt(), anyLong(), eq("testAlias3")))
+                .thenReturn(null);
+
+        AndroidKeyStoreSpi spi = new AndroidKeyStoreSpi();
+        spi.initForTesting(mKeystore2);
+
+        Enumeration<String> aliases = spi.engineAliases();
+        assertThat("Should have more elements before first.", aliases.hasMoreElements());
+        Assert.assertEquals(aliases.nextElement(), alias1);
+        assertThat("Should have more elements before second.", aliases.hasMoreElements());
+        Assert.assertEquals(aliases.nextElement(), alias2);
+        assertThat("Should have more elements before third.", aliases.hasMoreElements());
+        Assert.assertEquals(aliases.nextElement(), alias3);
+        assertThat("Should have no more elements after third.", !aliases.hasMoreElements());
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), eq(null));
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), eq("testAlias3"));
+        verifyNoMoreInteractions(mKeystore2);
+    }
+
+    @Test
+    public void testEngineAliasesCorrectlyListsEntriesInMultipleBatches() throws Exception {
+        final int numEntries = 2500;
+        KeyDescriptor[] kds = createKeyDescriptorsArray(numEntries);
+        when(mKeystore2.listBatch(anyInt(), anyLong(), eq(null)))
+                .thenReturn(Arrays.copyOfRange(kds, 0, 1000));
+        when(mKeystore2.listBatch(anyInt(), anyLong(), eq("alias-999")))
+                .thenReturn(Arrays.copyOfRange(kds, 1000, 2000));
+        when(mKeystore2.listBatch(anyInt(), anyLong(), eq("alias-1999")))
+                .thenReturn(Arrays.copyOfRange(kds, 2000, 2500));
+        when(mKeystore2.listBatch(anyInt(), anyLong(), eq("alias-2499")))
+                .thenReturn(null);
+
+        AndroidKeyStoreSpi spi = new AndroidKeyStoreSpi();
+        spi.initForTesting(mKeystore2);
+
+        assertAliasListsEqual(kds, spi.engineAliases());
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), eq(null));
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), eq("alias-999"));
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), eq("alias-1999"));
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), eq("alias-2499"));
+        verifyNoMoreInteractions(mKeystore2);
+    }
+
+    @Test
+    public void testEngineAliasesCorrectlyListsEntriesWhenNumEntriesIsExactlyOneBatchSize()
+            throws Exception {
+        final int numEntries = 1000;
+        KeyDescriptor[] kds = createKeyDescriptorsArray(numEntries);
+        when(mKeystore2.listBatch(anyInt(), anyLong(), eq(null)))
+                .thenReturn(kds);
+        when(mKeystore2.listBatch(anyInt(), anyLong(), eq("alias-999")))
+                .thenReturn(null);
+
+        AndroidKeyStoreSpi spi = new AndroidKeyStoreSpi();
+        spi.initForTesting(mKeystore2);
+
+        assertAliasListsEqual(kds, spi.engineAliases());
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), eq(null));
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), eq("alias-999"));
+        verifyNoMoreInteractions(mKeystore2);
+    }
+
+    @Test
+    public void testEngineAliasesCorrectlyListsEntriesWhenNumEntriesIsAMultiplyOfBatchSize()
+            throws Exception {
+        final int numEntries = 2000;
+        KeyDescriptor[] kds = createKeyDescriptorsArray(numEntries);
+        when(mKeystore2.listBatch(anyInt(), anyLong(), eq(null)))
+                .thenReturn(Arrays.copyOfRange(kds, 0, 1000));
+        when(mKeystore2.listBatch(anyInt(), anyLong(), eq("alias-999")))
+                .thenReturn(Arrays.copyOfRange(kds, 1000, 2000));
+        when(mKeystore2.listBatch(anyInt(), anyLong(), eq("alias-1999")))
+                .thenReturn(null);
+
+        AndroidKeyStoreSpi spi = new AndroidKeyStoreSpi();
+        spi.initForTesting(mKeystore2);
+
+        assertAliasListsEqual(kds, spi.engineAliases());
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), eq(null));
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), eq("alias-999"));
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), eq("alias-1999"));
+        verifyNoMoreInteractions(mKeystore2);
+    }
+
+    @Test
+    public void testEngineAliasesCorrectlyListsEntriesWhenReturningLessThanBatchSize()
+            throws Exception {
+        final int numEntries = 500;
+        KeyDescriptor[] kds = createKeyDescriptorsArray(numEntries);
+        when(mKeystore2.listBatch(anyInt(), anyLong(), eq(null)))
+                .thenReturn(kds);
+        when(mKeystore2.listBatch(anyInt(), anyLong(), eq("alias-499")))
+                .thenReturn(new KeyDescriptor[0]);
+
+        AndroidKeyStoreSpi spi = new AndroidKeyStoreSpi();
+        spi.initForTesting(mKeystore2);
+
+        assertAliasListsEqual(kds, spi.engineAliases());
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), eq(null));
+        verify(mKeystore2).listBatch(anyInt(), anyLong(), eq("alias-499"));
+        verifyNoMoreInteractions(mKeystore2);
     }
 
     @Mock