Notify DO after PO ownership transfer when they belong to affiliated users.

Test: Locally modified TransferDPC to print the callback after manually setting DO, PO and then transfering PO owner.
Bug: 72218214
Change-Id: I6c9af00bcfa703c81d9d5b35fe534c928991ff93
diff --git a/api/current.txt b/api/current.txt
index 0711ecd..5e91789 100644
--- a/api/current.txt
+++ b/api/current.txt
@@ -6348,6 +6348,7 @@
     method public void onReceive(android.content.Context, android.content.Intent);
     method public void onSecurityLogsAvailable(android.content.Context, android.content.Intent);
     method public void onSystemUpdatePending(android.content.Context, android.content.Intent, long);
+    method public void onTransferAffiliatedProfileOwnershipComplete(android.content.Context, android.os.UserHandle);
     method public void onTransferOwnershipComplete(android.content.Context, android.os.PersistableBundle);
     method public void onUserAdded(android.content.Context, android.content.Intent, android.os.UserHandle);
     method public void onUserRemoved(android.content.Context, android.content.Intent, android.os.UserHandle);
diff --git a/core/java/android/app/admin/DeviceAdminReceiver.java b/core/java/android/app/admin/DeviceAdminReceiver.java
index d65545d..28e845a 100644
--- a/core/java/android/app/admin/DeviceAdminReceiver.java
+++ b/core/java/android/app/admin/DeviceAdminReceiver.java
@@ -483,6 +483,16 @@
             "android.app.action.TRANSFER_OWNERSHIP_COMPLETE";
 
     /**
+     * Broadcast action: notify the device owner that the ownership of one of its affiliated
+     * profiles is transferred.
+     *
+     * @hide
+     */
+    @SdkConstant(SdkConstantType.BROADCAST_INTENT_ACTION)
+    public static final String ACTION_AFFILIATED_PROFILE_TRANSFER_OWNERSHIP_COMPLETE =
+            "android.app.action.AFFILIATED_PROFILE_TRANSFER_OWNERSHIP_COMPLETE";
+
+    /**
      * A {@link android.os.Parcelable} extra of type {@link android.os.PersistableBundle} that
      * allows a mobile device management application to pass data to the management application
      * instance after owner transfer.
@@ -993,6 +1003,26 @@
     }
 
     /**
+     * Called on the device owner when the ownership of one of its affiliated profiles is
+     * transferred.
+     *
+     * <p>This can be used when transferring both device and profile ownership when using
+     * work profile on a fully managed device. The process would look like this:
+     * <ol>
+     * <li>Transfer profile ownership</li>
+     * <li>The device owner gets notified with this callback</li>
+     * <li>Transfer device ownership</li>
+     * <li>Both profile and device ownerships have been transferred</li>
+     * </ol>
+     *
+     * @param context the running context as per {@link #onReceive}
+     * @param user the {@link UserHandle} of the affiliated user
+     * @see DevicePolicyManager#transferOwnership(ComponentName, ComponentName, PersistableBundle)
+     */
+    public void onTransferAffiliatedProfileOwnershipComplete(Context context, UserHandle user) {
+    }
+
+    /**
      * Intercept standard device administrator broadcasts.  Implementations
      * should not override this method; it is better to implement the
      * convenience callbacks for each action.
@@ -1064,6 +1094,9 @@
             PersistableBundle bundle =
                     intent.getParcelableExtra(EXTRA_TRANSFER_OWNERSHIP_ADMIN_EXTRAS_BUNDLE);
             onTransferOwnershipComplete(context, bundle);
+        } else if (ACTION_AFFILIATED_PROFILE_TRANSFER_OWNERSHIP_COMPLETE.equals(action)) {
+            onTransferAffiliatedProfileOwnershipComplete(context,
+                    intent.getParcelableExtra(Intent.EXTRA_USER));
         }
     }
 }
diff --git a/core/res/AndroidManifest.xml b/core/res/AndroidManifest.xml
index b046808..f1a9bd4 100644
--- a/core/res/AndroidManifest.xml
+++ b/core/res/AndroidManifest.xml
@@ -578,6 +578,7 @@
     <!-- Added in P -->
     <protected-broadcast android:name="android.app.action.PROFILE_OWNER_CHANGED" />
     <protected-broadcast android:name="android.app.action.TRANSFER_OWNERSHIP_COMPLETE" />
+    <protected-broadcast android:name="android.app.action.AFFILIATED_PROFILE_TRANSFER_OWNERSHIP_COMPLETE" />
     <protected-broadcast android:name="android.app.action.DATA_SHARING_RESTRICTION_CHANGED" />
 
     <!-- ====================================================================== -->
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index a45ca42..6e29196 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -12396,6 +12396,9 @@
                     sendProfileOwnerCommand(DeviceAdminReceiver.ACTION_TRANSFER_OWNERSHIP_COMPLETE,
                             getTransferOwnershipAdminExtras(bundle), callingUserId);
                     postTransfer(DevicePolicyManager.ACTION_PROFILE_OWNER_CHANGED, callingUserId);
+                    if (isUserAffiliatedWithDeviceLocked(callingUserId)) {
+                        notifyAffiliatedProfileTransferOwnershipComplete(callingUserId);
+                    }
                 } else if (isDeviceOwner(admin, callingUserId)) {
                     prepareTransfer(admin, target, bundle, callingUserId,
                             ADMIN_TYPE_DEVICE_OWNER);
@@ -12423,6 +12426,13 @@
         sendOwnerChangedBroadcast(broadcast, callingUserId);
     }
 
+    private void notifyAffiliatedProfileTransferOwnershipComplete(int callingUserId) {
+        final Bundle extras = new Bundle();
+        extras.putParcelable(Intent.EXTRA_USER, UserHandle.of(callingUserId));
+        sendDeviceOwnerCommand(
+                DeviceAdminReceiver.ACTION_AFFILIATED_PROFILE_TRANSFER_OWNERSHIP_COMPLETE, extras);
+    }
+
     /**
      * Transfers the profile owner for user with id profileOwnerUserId from admin to target.
      */