Merge "Make DnD check work profile contacts" into rvc-dev
diff --git a/core/java/android/provider/ContactsContract.java b/core/java/android/provider/ContactsContract.java
index 146387e..31e0a06 100644
--- a/core/java/android/provider/ContactsContract.java
+++ b/core/java/android/provider/ContactsContract.java
@@ -57,6 +57,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
+import java.util.List;
 
 /**
  * <p>
@@ -2210,6 +2211,28 @@
         public static InputStream openContactPhotoInputStream(ContentResolver cr, Uri contactUri) {
             return openContactPhotoInputStream(cr, contactUri, false);
         }
+
+        /**
+         * Creates and returns a corp lookup URI from the given enterprise lookup URI by removing
+         * {@link #ENTERPRISE_CONTACT_LOOKUP_PREFIX} from the key. Returns {@code null} if the given
+         * URI is not an enterprise lookup URI.
+         *
+         * @hide
+         */
+        @Nullable
+        public static Uri createCorpLookupUriFromEnterpriseLookupUri(
+                @NonNull Uri enterpriseLookupUri) {
+            final List<String> pathSegments = enterpriseLookupUri.getPathSegments();
+            if (pathSegments == null || pathSegments.size() <= 2) {
+                return null;
+            }
+            final String key = pathSegments.get(2);
+            if (TextUtils.isEmpty(key) || !key.startsWith(ENTERPRISE_CONTACT_LOOKUP_PREFIX)) {
+                return null;
+            }
+            final String actualKey = key.substring(ENTERPRISE_CONTACT_LOOKUP_PREFIX.length());
+            return Uri.withAppendedPath(Contacts.CONTENT_LOOKUP_URI, actualKey);
+        }
     }
 
     /**
diff --git a/services/core/java/com/android/server/notification/ValidateNotificationPeople.java b/services/core/java/com/android/server/notification/ValidateNotificationPeople.java
index 90fc59a..f59934f 100644
--- a/services/core/java/com/android/server/notification/ValidateNotificationPeople.java
+++ b/services/core/java/com/android/server/notification/ValidateNotificationPeople.java
@@ -19,6 +19,7 @@
 import android.annotation.Nullable;
 import android.app.Notification;
 import android.app.Person;
+import android.content.ContentProvider;
 import android.content.Context;
 import android.content.pm.PackageManager;
 import android.database.ContentObserver;
@@ -28,6 +29,7 @@
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.UserHandle;
+import android.os.UserManager;
 import android.provider.ContactsContract;
 import android.provider.ContactsContract.Contacts;
 import android.provider.Settings;
@@ -38,6 +40,8 @@
 import android.util.LruCache;
 import android.util.Slog;
 
+import com.android.internal.annotations.VisibleForTesting;
+
 import libcore.util.EmptyArray;
 
 import java.util.ArrayList;
@@ -392,26 +396,57 @@
         return searchContacts(context, numberUri);
     }
 
-    private LookupResult searchContacts(Context context, Uri lookupUri) {
+    @VisibleForTesting
+    LookupResult searchContacts(Context context, Uri lookupUri) {
         LookupResult lookupResult = new LookupResult();
-        Cursor c = null;
-        try {
-            c = context.getContentResolver().query(lookupUri, LOOKUP_PROJECTION, null, null, null);
+        final Uri corpLookupUri =
+                ContactsContract.Contacts.createCorpLookupUriFromEnterpriseLookupUri(lookupUri);
+        if (corpLookupUri == null) {
+            addContacts(lookupResult, context, lookupUri);
+        } else {
+            addWorkContacts(lookupResult, context, corpLookupUri);
+        }
+        return lookupResult;
+    }
+
+    private void addWorkContacts(LookupResult lookupResult, Context context, Uri corpLookupUri) {
+        final int workUserId = findWorkUserId(context);
+        if (workUserId == -1) {
+            Slog.w(TAG, "Work profile user ID not found for work contact: " + corpLookupUri);
+            return;
+        }
+        final Uri corpLookupUriWithUserId =
+                ContentProvider.maybeAddUserId(corpLookupUri, workUserId);
+        addContacts(lookupResult, context, corpLookupUriWithUserId);
+    }
+
+    /** Returns the user ID of the managed profile or -1 if none is found. */
+    private int findWorkUserId(Context context) {
+        final UserManager userManager = context.getSystemService(UserManager.class);
+        final int[] profileIds =
+                userManager.getProfileIds(context.getUserId(), /* enabledOnly= */ true);
+        for (int profileId : profileIds) {
+            if (userManager.isManagedProfile(profileId)) {
+                return profileId;
+            }
+        }
+        return -1;
+    }
+
+    /** Modifies the given lookup result to add contacts found at the given URI. */
+    private void addContacts(LookupResult lookupResult, Context context, Uri uri) {
+        try (Cursor c = context.getContentResolver().query(
+                uri, LOOKUP_PROJECTION, null, null, null)) {
             if (c == null) {
                 Slog.w(TAG, "Null cursor from contacts query.");
-                return lookupResult;
+                return;
             }
             while (c.moveToNext()) {
                 lookupResult.mergeContact(c);
             }
         } catch (Throwable t) {
             Slog.w(TAG, "Problem getting content resolver or performing contacts query.", t);
-        } finally {
-            if (c != null) {
-                c.close();
-            }
         }
-        return lookupResult;
     }
 
     private static class LookupResult {
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/ValidateNotificationPeopleTest.java b/services/tests/uiservicestests/src/com/android/server/notification/ValidateNotificationPeopleTest.java
index b1ee120..0bf105d 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/ValidateNotificationPeopleTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/ValidateNotificationPeopleTest.java
@@ -16,11 +16,23 @@
 package com.android.server.notification;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.isNull;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import android.app.Notification;
 import android.app.Person;
+import android.content.ContentProvider;
+import android.content.ContentResolver;
+import android.content.Context;
+import android.net.Uri;
 import android.os.Bundle;
+import android.os.UserManager;
+import android.provider.ContactsContract;
 import android.test.suitebuilder.annotation.SmallTest;
 import android.text.SpannableString;
 
@@ -30,6 +42,7 @@
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -169,6 +182,64 @@
         assertStringArrayEquals("testPeopleArrayList", expected, result);
     }
 
+    @Test
+    public void testSearchContacts_workContact_queriesWorkContactProvider()
+            throws Exception {
+        final int personalUserId = 0;
+        final int workUserId = 12;
+        final int contactId = 12345;
+        final Context mockContext = mock(Context.class);
+        when(mockContext.getUserId()).thenReturn(personalUserId);
+        final UserManager mockUserManager = mock(UserManager.class);
+        when(mockContext.getSystemService(UserManager.class)).thenReturn(mockUserManager);
+        when(mockUserManager.getProfileIds(personalUserId, /* enabledOnly= */ true))
+                .thenReturn(new int[] {personalUserId, workUserId});
+        when(mockUserManager.isManagedProfile(workUserId)).thenReturn(true);
+        final ContentResolver mockContentResolver = mock(ContentResolver.class);
+        when(mockContext.getContentResolver()).thenReturn(mockContentResolver);
+        final Uri lookupUri = Uri.withAppendedPath(
+                ContactsContract.Contacts.CONTENT_LOOKUP_URI,
+                ContactsContract.Contacts.ENTERPRISE_CONTACT_LOOKUP_PREFIX + contactId);
+
+        new ValidateNotificationPeople().searchContacts(mockContext, lookupUri);
+
+        ArgumentCaptor<Uri> queryUri = ArgumentCaptor.forClass(Uri.class);
+        verify(mockContentResolver).query(
+                queryUri.capture(),
+                any(),
+                /* selection= */ isNull(),
+                /* selectionArgs= */ isNull(),
+                /* sortOrder= */ isNull());
+        assertEquals(workUserId, ContentProvider.getUserIdFromUri(queryUri.getValue()));
+    }
+
+    @Test
+    public void testSearchContacts_personalContact_queriesPersonalContactProvider()
+            throws Exception {
+        final int personalUserId = 0;
+        final int workUserId = 12;
+        final int contactId = 12345;
+        final Context mockContext = mock(Context.class);
+        when(mockContext.getUserId()).thenReturn(personalUserId);
+        final UserManager mockUserManager = mock(UserManager.class);
+        when(mockContext.getSystemService(UserManager.class)).thenReturn(mockUserManager);
+        final ContentResolver mockContentResolver = mock(ContentResolver.class);
+        when(mockContext.getContentResolver()).thenReturn(mockContentResolver);
+        final Uri lookupUri = Uri.withAppendedPath(
+                ContactsContract.Contacts.CONTENT_LOOKUP_URI, String.valueOf(contactId));
+
+        new ValidateNotificationPeople().searchContacts(mockContext, lookupUri);
+
+        ArgumentCaptor<Uri> queryUri = ArgumentCaptor.forClass(Uri.class);
+        verify(mockContentResolver).query(
+                queryUri.capture(),
+                any(),
+                /* selection= */ isNull(),
+                /* selectionArgs= */ isNull(),
+                /* sortOrder= */ isNull());
+        assertFalse(ContentProvider.uriHasUserId(queryUri.getValue()));
+    }
+
     private void assertStringArrayEquals(String message, String[] expected, String[] result) {
         String expectedString = Arrays.toString(expected);
         String resultString = Arrays.toString(result);