Fix validation crash for nextPageToken.

We are caching nextPageToken in query(), and validate it in
getNextPage(). But when it reach to the end, the token will be changed
to 0. And it will crash if user is trying to fetch next page since 0
is not cached and won't passed the validation.

This change fix the issue. Without it, users maybe crashed if they
are trying to retrieve all documents in all result pages.

Bug: 187972715
Test: AppSearchSessionUnitTest
Change-Id: Ia95feb52f2a37348e11afdf0b320c61bfce22d40

3 files changed, 150 insertions, 1 deletions

diff --git a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java
index 830e76c62279..fd2ff0bc6b7e 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java
@@ -145,6 +145,9 @@ import java.util.concurrent.locks.ReentrantReadWriteLock;
 public final class AppSearchImpl implements Closeable {
     private static final String TAG = "AppSearchImpl";
 
+    /** A value 0 means that there're no more pages in the search results. */
+    private static final long EMPTY_PAGE_TOKEN = 0;
+
     @VisibleForTesting static final int CHECK_OPTIMIZE_INTERVAL = 100;
 
     private final ReadWriteLock mReadWriteLock = new ReentrantReadWriteLock();
@@ -1135,6 +1138,16 @@ public final class AppSearchImpl implements Closeable {
                     searchResultProto.getResultsCount(),
                     searchResultProto);
             checkSuccess(searchResultProto.getStatus());
+            if (nextPageToken != EMPTY_PAGE_TOKEN
+                    && searchResultProto.getNextPageToken() == EMPTY_PAGE_TOKEN) {
+                // At this point, we're guaranteed that this nextPageToken exists for this package,
+                // otherwise checkNextPageToken would've thrown an exception.
+                // Since the new token is 0, this is the last page. We should remove the old token
+                // from our cache since it no longer refers to this query.
+                synchronized (mNextPageTokensLocked) {
+                    mNextPageTokensLocked.get(packageName).remove(nextPageToken);
+                }
+            }
             return rewriteSearchResultProto(searchResultProto, mSchemaMapLocked);
         } finally {
             mReadWriteLock.readLock().unlock();
@@ -2056,6 +2069,10 @@ public final class AppSearchImpl implements Closeable {
     }
 
     private void addNextPageToken(String packageName, long nextPageToken) {
+        if (nextPageToken == EMPTY_PAGE_TOKEN) {
+            // There is no more pages. No need to add it.
+            return;
+        }
         synchronized (mNextPageTokensLocked) {
             Set<Long> tokens = mNextPageTokensLocked.get(packageName);
             if (tokens == null) {
@@ -2068,6 +2085,11 @@ public final class AppSearchImpl implements Closeable {
 
     private void checkNextPageToken(String packageName, long nextPageToken)
             throws AppSearchException {
+        if (nextPageToken == EMPTY_PAGE_TOKEN) {
+            // Swallow the check for empty page token, token = 0 means there is no more page and it
+            // won't return anything from Icing.
+            return;
+        }
         synchronized (mNextPageTokensLocked) {
             Set<Long> nextPageTokens = mNextPageTokensLocked.get(packageName);
             if (nextPageTokens == null || !nextPageTokens.contains(nextPageToken)) {
diff --git a/apex/appsearch/synced_jetpack_changeid.txt b/apex/appsearch/synced_jetpack_changeid.txt
index 65551074f9c0..a81d7d8022b2 100644
--- a/apex/appsearch/synced_jetpack_changeid.txt
+++ b/apex/appsearch/synced_jetpack_changeid.txt
@@ -1 +1 @@
-c7387d9b58726a23a0608a9365fb3a1b57b7b8a1
+Ie04f1ecc033faae8085afcb51eb9e40a298998d5
diff --git a/core/tests/coretests/src/android/app/appsearch/AppSearchSessionUnitTest.java b/core/tests/coretests/src/android/app/appsearch/AppSearchSessionUnitTest.java
index d51004c08585..07e4333e74b1 100644
--- a/core/tests/coretests/src/android/app/appsearch/AppSearchSessionUnitTest.java
+++ b/core/tests/coretests/src/android/app/appsearch/AppSearchSessionUnitTest.java
@@ -16,6 +16,8 @@
 
 package android.app.appsearch;
 
+import static android.app.appsearch.SearchSpec.TERM_MATCH_PREFIX;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.testng.Assert.expectThrows;
@@ -30,6 +32,8 @@ import com.android.server.appsearch.testing.AppSearchEmail;
 import org.junit.Before;
 import org.junit.Test;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
@@ -100,4 +104,127 @@ public class AppSearchSessionUnitTest {
                 .isEqualTo(AppSearchResult.RESULT_INTERNAL_ERROR);
         assertThat(appSearchException.getMessage()).startsWith("NullPointerException");
     }
+
+    @Test
+    public void testGetEmptyNextPage() throws Exception {
+        // Set the schema.
+        CompletableFuture<AppSearchResult<SetSchemaResponse>> schemaFuture =
+                new CompletableFuture<>();
+        mSearchSession.setSchema(
+                new SetSchemaRequest.Builder()
+                        .addSchemas(new AppSearchSchema.Builder("schema1").build())
+                        .setForceOverride(true).build(),
+                mExecutor, mExecutor, schemaFuture::complete);
+        schemaFuture.get().getResultValue();
+
+        // Create a document and index it.
+        GenericDocument document1 = new GenericDocument.Builder<>("namespace", "id1",
+                "schema1").build();
+        CompletableFuture<AppSearchBatchResult<String, Void>> putDocumentsFuture =
+                new CompletableFuture<>();
+        mSearchSession.put(
+                new PutDocumentsRequest.Builder().addGenericDocuments(document1).build(),
+                mExecutor, new BatchResultCallback<String, Void>() {
+                    @Override
+                    public void onResult(AppSearchBatchResult<String, Void> result) {
+                        putDocumentsFuture.complete(result);
+                    }
+
+                    @Override
+                    public void onSystemError(Throwable throwable) {
+                        putDocumentsFuture.completeExceptionally(throwable);
+                    }
+                });
+        putDocumentsFuture.get();
+
+        // Search and get the first page.
+        SearchSpec searchSpec = new SearchSpec.Builder()
+                .setTermMatch(TERM_MATCH_PREFIX)
+                .setResultCountPerPage(1)
+                .build();
+        SearchResults searchResults = mSearchSession.search("", searchSpec);
+
+        CompletableFuture<AppSearchResult<List<SearchResult>>> getNextPageFuture =
+                new CompletableFuture<>();
+        searchResults.getNextPage(mExecutor, getNextPageFuture::complete);
+        List<SearchResult> results = getNextPageFuture.get().getResultValue();
+        assertThat(results).hasSize(1);
+        assertThat(results.get(0).getGenericDocument()).isEqualTo(document1);
+
+        // We get all documents, and it shouldn't fail if we keep calling getNextPage().
+        getNextPageFuture = new CompletableFuture<>();
+        searchResults.getNextPage(mExecutor, getNextPageFuture::complete);
+        results = getNextPageFuture.get().getResultValue();
+        assertThat(results).isEmpty();
+    }
+
+    @Test
+    public void testGetEmptyNextPage_multiPages() throws Exception {
+        // Set the schema.
+        CompletableFuture<AppSearchResult<SetSchemaResponse>> schemaFuture =
+                new CompletableFuture<>();
+        mSearchSession.setSchema(
+                new SetSchemaRequest.Builder()
+                        .addSchemas(new AppSearchSchema.Builder("schema1").build())
+                        .setForceOverride(true).build(),
+                mExecutor, mExecutor, schemaFuture::complete);
+        schemaFuture.get().getResultValue();
+
+        // Create a document and insert 3 package1 documents
+        GenericDocument document1 = new GenericDocument.Builder<>("namespace", "id1",
+                "schema1").build();
+        GenericDocument document2 = new GenericDocument.Builder<>("namespace", "id2",
+                "schema1").build();
+        GenericDocument document3 = new GenericDocument.Builder<>("namespace", "id3",
+                "schema1").build();
+        CompletableFuture<AppSearchBatchResult<String, Void>> putDocumentsFuture =
+                new CompletableFuture<>();
+        mSearchSession.put(
+                new PutDocumentsRequest.Builder()
+                        .addGenericDocuments(document1, document2, document3).build(),
+                mExecutor, new BatchResultCallback<String, Void>() {
+                    @Override
+                    public void onResult(AppSearchBatchResult<String, Void> result) {
+                        putDocumentsFuture.complete(result);
+                    }
+
+                    @Override
+                    public void onSystemError(Throwable throwable) {
+                        putDocumentsFuture.completeExceptionally(throwable);
+                    }
+                });
+        putDocumentsFuture.get();
+
+        // Search for only 2 result per page
+        SearchSpec searchSpec = new SearchSpec.Builder()
+                .setTermMatch(TERM_MATCH_PREFIX)
+                .setResultCountPerPage(2)
+                .build();
+        SearchResults searchResults = mSearchSession.search("", searchSpec);
+
+        // Get the first page, it contains 2 results.
+        List<GenericDocument> outDocs = new ArrayList<>();
+        CompletableFuture<AppSearchResult<List<SearchResult>>> getNextPageFuture =
+                new CompletableFuture<>();
+        searchResults.getNextPage(mExecutor, getNextPageFuture::complete);
+        List<SearchResult> results = getNextPageFuture.get().getResultValue();
+        assertThat(results).hasSize(2);
+        outDocs.add(results.get(0).getGenericDocument());
+        outDocs.add(results.get(1).getGenericDocument());
+
+        // Get the second page, it contains only 1 result.
+        getNextPageFuture = new CompletableFuture<>();
+        searchResults.getNextPage(mExecutor, getNextPageFuture::complete);
+        results = getNextPageFuture.get().getResultValue();
+        assertThat(results).hasSize(1);
+        outDocs.add(results.get(0).getGenericDocument());
+
+        assertThat(outDocs).containsExactly(document1, document2, document3);
+
+        // We get all documents, and it shouldn't fail if we keep calling getNextPage().
+        getNextPageFuture = new CompletableFuture<>();
+        searchResults.getNextPage(mExecutor, getNextPageFuture::complete);
+        results = getNextPageFuture.get().getResultValue();
+        assertThat(results).isEmpty();
+    }
 }