From 040121465b6c4ab811b9bede437ec261d5285029 Mon Sep 17 00:00:00 2001
From: Simranjit Kohli <simranjit@google.com>
Date: Thu, 10 Aug 2023 05:27:02 -0700
Subject: [PATCH] [Autofill PCC] Fix authentication logic for prefer PCC

Fix the authentication filling logic when prefer PCC is enabled.
Previous internal logic could create multiple datasets after authenticaiton.
This problem was visible if the PCC didn't detect any of the filling fields.
The bug is fixed by merging the datasets with the detected fields, as to
create a single dataset.

Bug: 293313411
Fixes: 291194194
Test: atest CtsAutoFillServiceTestCases
Change-Id: I5dd50c6639a96a046213e00fddf0b15c8bb3d5e3
---
 .../com/android/server/autofill/Session.java  | 39 +++++++++++++++++--
 1 file changed, 35 insertions(+), 4 deletions(-)

diff --git a/services/autofill/java/com/android/server/autofill/Session.java b/services/autofill/java/com/android/server/autofill/Session.java
index e19ef1074ebf8..65e1916c27de4 100644
--- a/services/autofill/java/com/android/server/autofill/Session.java
+++ b/services/autofill/java/com/android/server/autofill/Session.java
@@ -134,6 +134,7 @@ import android.service.autofill.AutofillService;
 import android.service.autofill.CompositeUserData;
 import android.service.autofill.Dataset;
 import android.service.autofill.Dataset.DatasetEligibleReason;
+import android.service.autofill.Field;
 import android.service.autofill.FieldClassification;
 import android.service.autofill.FieldClassification.Match;
 import android.service.autofill.FieldClassificationUserData;
@@ -2578,10 +2579,7 @@ final class Session implements RemoteFillService.FillServiceCallbacks, ViewState
                     if (sDebug) Slog.d(TAG,  "Updating client state from auth dataset");
                     mClientState = newClientState;
                 }
-                Dataset dataset = (Dataset) result;
-                FillResponse temp = new FillResponse.Builder().addDataset(dataset).build();
-                temp = getEffectiveFillResponse(temp);
-                dataset = temp.getDatasets().get(0);
+                Dataset dataset = getEffectiveDatasetForAuthentication((Dataset) result);
                 final Dataset oldDataset = authenticatedResponse.getDatasets().get(datasetIdx);
                 if (!isAuthResultDatasetEphemeral(oldDataset, data)) {
                     authenticatedResponse.getDatasets().set(datasetIdx, dataset);
@@ -2607,6 +2605,39 @@ final class Session implements RemoteFillService.FillServiceCallbacks, ViewState
         }
     }
 
+    Dataset getEffectiveDatasetForAuthentication(Dataset authenticatedDataset) {
+        FillResponse response = new FillResponse.Builder().addDataset(authenticatedDataset).build();
+        response = getEffectiveFillResponse(response);
+        if (DBG) {
+            Slog.d(TAG, "DBG: authenticated effective response: " + response);
+        }
+        if (response == null || response.getDatasets().size() == 0) {
+            Log.wtf(TAG, "No datasets in fill response on authentication. response = "
+                    + (response == null ? "null" : response.toString()));
+            return authenticatedDataset;
+        }
+        List<Dataset> datasets = response.getDatasets();
+        Dataset result = response.getDatasets().get(0);
+        if (datasets.size() > 1) {
+            Dataset.Builder builder = new Dataset.Builder();
+            for (Dataset dataset : datasets) {
+                if (!dataset.getFieldIds().isEmpty()) {
+                    for (int i = 0; i < dataset.getFieldIds().size(); i++) {
+                        builder.setField(dataset.getFieldIds().get(i),
+                                new Field.Builder().setValue(dataset.getFieldValues().get(i))
+                                        .build());
+                    }
+                }
+            }
+            result = builder.setId(authenticatedDataset.getId()).build();
+        }
+
+        if (DBG) {
+            Slog.d(TAG, "DBG: authenticated effective dataset after auth: " + result);
+        }
+        return result;
+    }
+
     /**
      * Returns whether the dataset returned from the authentication result is ephemeral or not.
      * See {@link AutofillManager#EXTRA_AUTHENTICATION_RESULT_EPHEMERAL_DATASET} for more