Sets AutofillClient on attachBaseContext()

When an activity overrides attachBaseContext(), the attached context must share
the same AutofillCient, otherwise it will break autofill in the activity.

Test: cts-tradefed run commandAndExit cts-dev -m CtsAutoFillServiceTestCases -t android.autofillservice.cts.AttachedContextActivityTest
Test: cts-tradefed run commandAndExit cts-dev -m CtsAutoFillServiceTestCases

Fixes: 67350639

Change-Id: I6b0c4c2fab5fa68eb6ee0714e107ca0816efb2cb

5 files changed, 61 insertions, 4 deletions

diff --git a/core/java/android/app/Activity.java b/core/java/android/app/Activity.java
index 785a8f7118ed..e14b7382e827 100644
--- a/core/java/android/app/Activity.java
+++ b/core/java/android/app/Activity.java
@@ -114,6 +114,7 @@ import android.view.WindowManager;
 import android.view.WindowManagerGlobal;
 import android.view.accessibility.AccessibilityEvent;
 import android.view.autofill.AutofillManager;
+import android.view.autofill.AutofillManager.AutofillClient;
 import android.view.autofill.AutofillPopupWindow;
 import android.view.autofill.IAutofillWindowPresenter;
 import android.widget.AdapterView;
@@ -947,6 +948,18 @@ public class Activity extends ContextThemeWrapper
         return mAutofillManager;
     }
 
+    @Override
+    protected void attachBaseContext(Context newBase) {
+        super.attachBaseContext(newBase);
+        newBase.setAutofillClient(this);
+    }
+
+    /** @hide */
+    @Override
+    public final AutofillClient getAutofillClient() {
+        return this;
+    }
+
     /**
      * Called when the activity is starting.  This is where most initialization
      * should go: calling {@link #setContentView(int)} to inflate the
diff --git a/core/java/android/app/ContextImpl.java b/core/java/android/app/ContextImpl.java
index c48be7705706..5f3432264ca0 100644
--- a/core/java/android/app/ContextImpl.java
+++ b/core/java/android/app/ContextImpl.java
@@ -74,6 +74,7 @@ import android.util.Log;
 import android.util.Slog;
 import android.view.Display;
 import android.view.DisplayAdjustments;
+import android.view.autofill.AutofillManager.AutofillClient;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.util.Preconditions;
@@ -185,6 +186,8 @@ class ContextImpl extends Context {
     // The name of the split this Context is representing. May be null.
     private @Nullable String mSplitName = null;
 
+    private AutofillClient mAutofillClient = null;
+
     private final Object mSync = new Object();
 
     @GuardedBy("mSync")
@@ -2225,6 +2228,18 @@ class ContextImpl extends Context {
         return mUser.getIdentifier();
     }
 
+    /** @hide */
+    @Override
+    public AutofillClient getAutofillClient() {
+        return mAutofillClient;
+    }
+
+    /** @hide */
+    @Override
+    public void setAutofillClient(AutofillClient client) {
+        mAutofillClient = client;
+    }
+
     static ContextImpl createSystemContext(ActivityThread mainThread) {
         LoadedApk packageInfo = new LoadedApk(mainThread);
         ContextImpl context = new ContextImpl(null, mainThread, packageInfo, null, null, null, 0,
diff --git a/core/java/android/content/Context.java b/core/java/android/content/Context.java
index 2d8249acb5bf..792e8f37f467 100644
--- a/core/java/android/content/Context.java
+++ b/core/java/android/content/Context.java
@@ -64,6 +64,7 @@ import android.view.DisplayAdjustments;
 import android.view.View;
 import android.view.ViewDebug;
 import android.view.WindowManager;
+import android.view.autofill.AutofillManager.AutofillClient;
 import android.view.textclassifier.TextClassificationManager;
 
 import java.io.File;
@@ -4765,6 +4766,19 @@ public abstract class Context {
     }
 
     /**
+     * @hide
+     */
+    public AutofillClient getAutofillClient() {
+        return null;
+    }
+
+    /**
+     * @hide
+     */
+    public void setAutofillClient(AutofillClient client) {
+    }
+
+    /**
      * Throws an exception if the Context is using system resources,
      * which are non-runtime-overlay-themable and may show inconsistent UI.
      * @hide
diff --git a/core/java/android/content/ContextWrapper.java b/core/java/android/content/ContextWrapper.java
index a9fd58bc950c..85acdc6b8101 100644
--- a/core/java/android/content/ContextWrapper.java
+++ b/core/java/android/content/ContextWrapper.java
@@ -37,6 +37,7 @@ import android.os.Looper;
 import android.os.UserHandle;
 import android.view.Display;
 import android.view.DisplayAdjustments;
+import android.view.autofill.AutofillManager.AutofillClient;
 
 import java.io.File;
 import java.io.FileInputStream;
@@ -967,7 +968,24 @@ public class ContextWrapper extends Context {
     /**
      * @hide
      */
+    @Override
     public int getNextAutofillId() {
         return mBase.getNextAutofillId();
     }
+
+    /**
+     * @hide
+     */
+    @Override
+    public AutofillClient getAutofillClient() {
+        return mBase.getAutofillClient();
+    }
+
+    /**
+     * @hide
+     */
+    @Override
+    public void setAutofillClient(AutofillClient client) {
+        mBase.setAutofillClient(client);
+    }
 }
diff --git a/core/java/android/view/autofill/AutofillManager.java b/core/java/android/view/autofill/AutofillManager.java
index e906a1fa7807..4fb2a99af575 100644
--- a/core/java/android/view/autofill/AutofillManager.java
+++ b/core/java/android/view/autofill/AutofillManager.java
@@ -937,10 +937,7 @@ public final class AutofillManager {
     }
 
     private AutofillClient getClientLocked() {
-        if (mContext instanceof AutofillClient) {
-            return (AutofillClient) mContext;
-        }
-        return null;
+        return mContext.getAutofillClient();
     }
 
     /** @hide */