Open up new API in StagingManager to get information about staged APEX

These APIs will be later used by PackageManagerNative service to open up
the information to native clients.

The current implementation is a bit wasteful since everytime the client
makes a request, we ask the ApexManager to fetch us information about
the staged sessions N times, where N is the number of staged sessions
that are marked ready.

We can optimize it by caching the result of getStagedApexInfos(session)
method. This will be done in a separate CL.

Bug: 187444679
Test: atest StagingManagerTest
Change-Id: I76803382db26a029f6d9cabb0d9ce9a2bf9c8ea1
Merged-In: I76803382db26a029f6d9cabb0d9ce9a2bf9c8ea1
(cherry picked from commit c8132a9437)

services/core/java/com/android/server/pm/StagingManager.java

@@ -17,6 +17,7 @@
 package com.android.server.pm;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.apex.ApexInfo;
 import android.apex.ApexInfoList;
 import android.apex.ApexSessionInfo;
@@ -52,6 +53,7 @@ import android.os.SystemProperties;
 import android.os.Trace;
 import android.os.UserHandle;
 import android.text.TextUtils;
+import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.IntArray;
 import android.util.Slog;
@@ -80,7 +82,9 @@ import java.io.File;
 import java.io.FileReader;
 import java.io.FileWriter;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
@@ -1134,6 +1138,86 @@ public class StagingManager {
         mPreRebootVerificationHandler.notifyPreRebootVerification_Apk_Complete(session);
     }
 
+    /**
+     * Returns ApexInfo about APEX contained inside the session as a {@code Map<String, ApexInfo>},
+     * where the key of the map is the module name of the ApexInfo.
+     *
+     * Returns an empty map if there is any error.
+     */
+    @VisibleForTesting
+    @NonNull
+    Map<String, ApexInfo> getStagedApexInfos(@NonNull StagedSession session) {
+        Preconditions.checkArgument(session != null, "Session is null");
+        Preconditions.checkArgument(!session.hasParentSessionId(),
+                session.sessionId() + " session has parent session");
+        Preconditions.checkArgument(session.containsApexSession(),
+                session.sessionId() + " session does not contain apex");
+
+        // Even if caller calls this method on ready session, the session could be abandoned
+        // right after this method is called.
+        if (!session.isSessionReady() || session.isDestroyed()) {
+            return Collections.emptyMap();
+        }
+
+        ApexSessionParams params = new ApexSessionParams();
+        params.sessionId = session.sessionId();
+        final IntArray childSessionIds = new IntArray();
+        if (session.isMultiPackage()) {
+            for (StagedSession s : session.getChildSessions()) {
+                if (s.isApexSession()) {
+                    childSessionIds.add(s.sessionId());
+                }
+            }
+        }
+        params.childSessionIds = childSessionIds.toArray();
+
+        ApexInfo[] infos = mApexManager.getStagedApexInfos(params);
+        Map<String, ApexInfo> result = new ArrayMap<>();
+        for (ApexInfo info : infos) {
+            result.put(info.moduleName, info);
+        }
+        return result;
+    }
+
+    /**
+     * Returns apex module names of all packages that are staged ready
+     */
+    List<String> getStagedApexModuleNames() {
+        List<String> result = new ArrayList<>();
+        synchronized (mStagedSessions) {
+            for (int i = 0; i < mStagedSessions.size(); i++) {
+                final StagedSession session = mStagedSessions.valueAt(i);
+                if (!session.isSessionReady() || session.isDestroyed()
+                        || session.hasParentSessionId() || !session.containsApexSession()) {
+                    continue;
+                }
+                result.addAll(getStagedApexInfos(session).keySet());
+            }
+        }
+        return result;
+    }
+
+    /**
+     * Returns ApexInfo of the {@code moduleInfo} provided if it is staged, otherwise returns null.
+     */
+    @Nullable
+    ApexInfo getStagedApexInfo(String moduleName) {
+        synchronized (mStagedSessions) {
+            for (int i = 0; i < mStagedSessions.size(); i++) {
+                final StagedSession session = mStagedSessions.valueAt(i);
+                if (!session.isSessionReady() || session.isDestroyed()
+                        || session.hasParentSessionId() || !session.containsApexSession()) {
+                    continue;
+                }
+                ApexInfo result = getStagedApexInfos(session).get(moduleName);
+                if (result != null) {
+                    return result;
+                }
+            }
+        }
+        return null;
+    }
+
     private final class PreRebootVerificationHandler extends Handler {
         // Hold sessions before handler gets ready to do the verification.
         private List<StagedSession> mPendingSessions;

services/tests/mockingservicestests/src/com/android/server/pm/StagingManagerTest.java

@@ -18,6 +18,7 @@ package com.android.server.pm;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.junit.Assert.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
@@ -28,9 +29,10 @@ import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
-import static org.testng.Assert.assertThrows;
 
+import android.apex.ApexInfo;
 import android.apex.ApexSessionInfo;
+import android.apex.ApexSessionParams;
 import android.content.Context;
 import android.content.IntentSender;
 import android.content.pm.PackageInstaller;
@@ -39,6 +41,7 @@ import android.content.pm.PackageInstaller.SessionInfo.StagedSessionErrorCode;
 import android.os.SystemProperties;
 import android.os.storage.IStorageManager;
 import android.platform.test.annotations.Presubmit;
+import android.util.IntArray;
 import android.util.SparseArray;
 
 import com.android.dx.mockito.inline.extended.ExtendedMockito;
@@ -53,15 +56,19 @@ import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.mockito.MockitoSession;
+import org.mockito.invocation.InvocationOnMock;
 import org.mockito.quality.Strictness;
+import org.mockito.stubbing.Answer;
 
 import java.io.File;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 import java.util.function.Predicate;
 
 @Presubmit
@@ -519,6 +526,174 @@ public class StagingManagerTest {
         assertThat(mStagingManager.getSessionIdByPackageName("com.bar")).isEqualTo(-1);
     }
 
+    @Test
+    public void getStagedApexInfos_validatePreConditions() throws Exception {
+        // Invalid session: null session
+        {
+            // Call and verify
+            IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class,
+                    () -> mStagingManager.getStagedApexInfos(null));
+            assertThat(thrown).hasMessageThat().contains("Session is null");
+        }
+        // Invalid session: has parent
+        {
+            FakeStagedSession session = new FakeStagedSession(241);
+            session.setParentSessionId(239);
+            session.setSessionReady();
+            // Call and verify
+            IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class,
+                    () -> mStagingManager.getStagedApexInfos(session));
+            assertThat(thrown).hasMessageThat().contains("241 session has parent");
+        }
+
+        // Invalid session: does not contain apex
+        {
+            FakeStagedSession session = new FakeStagedSession(241);
+            session.setSessionReady();
+            // Call and verify
+            IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class,
+                    () -> mStagingManager.getStagedApexInfos(session));
+            assertThat(thrown).hasMessageThat().contains("241 session does not contain apex");
+        }
+        // Invalid session: not ready
+        {
+            FakeStagedSession session = new FakeStagedSession(239);
+            session.setIsApex(true);
+            // Call and verify
+            Map<String, ApexInfo> result = mStagingManager.getStagedApexInfos(session);
+            assertThat(result).isEmpty();
+        }
+        // Invalid session: destroyed
+        {
+            FakeStagedSession session = new FakeStagedSession(240);
+            session.setSessionReady();
+            session.setIsApex(true);
+            session.setDestroyed(true);
+            // Call and verify
+            Map<String, ApexInfo> result = mStagingManager.getStagedApexInfos(session);
+            assertThat(result).isEmpty();
+        }
+    }
+
+    private ApexInfo[] getFakeApexInfo(List<String> moduleNames) {
+        List<ApexInfo> result = new ArrayList<>();
+        for (String moduleName : moduleNames) {
+            ApexInfo info = new ApexInfo();
+            info.moduleName = moduleName;
+            result.add(info);
+        }
+        return result.toArray(new ApexInfo[0]);
+    }
+
+    @Test
+    public void getStagedApexInfos_nonParentSession() throws Exception {
+        FakeStagedSession validSession = new FakeStagedSession(239);
+        validSession.setIsApex(true);
+        validSession.setSessionReady();
+        ApexInfo[] fakeApexInfos = getFakeApexInfo(Arrays.asList("module1"));
+        when(mApexManager.getStagedApexInfos(any())).thenReturn(fakeApexInfos);
+
+        // Call and verify
+        Map<String, ApexInfo> result = mStagingManager.getStagedApexInfos(validSession);
+        assertThat(result).containsExactly(fakeApexInfos[0].moduleName, fakeApexInfos[0]);
+
+        ArgumentCaptor<ApexSessionParams> argumentCaptor =
+                ArgumentCaptor.forClass(ApexSessionParams.class);
+        verify(mApexManager, times(1)).getStagedApexInfos(argumentCaptor.capture());
+        ApexSessionParams params = argumentCaptor.getValue();
+        assertThat(params.sessionId).isEqualTo(239);
+    }
+
+    @Test
+    public void getStagedApexInfos_parentSession() throws Exception {
+        FakeStagedSession childSession1 = new FakeStagedSession(201);
+        childSession1.setIsApex(true);
+        FakeStagedSession childSession2 = new FakeStagedSession(202);
+        childSession2.setIsApex(true);
+        FakeStagedSession nonApexChild = new FakeStagedSession(203);
+        FakeStagedSession parentSession = new FakeStagedSession(239,
+                Arrays.asList(childSession1, childSession2, nonApexChild));
+        parentSession.setSessionReady();
+        ApexInfo[] fakeApexInfos = getFakeApexInfo(Arrays.asList("module1", "module2"));
+        when(mApexManager.getStagedApexInfos(any())).thenReturn(fakeApexInfos);
+
+        // Call and verify
+        Map<String, ApexInfo> result = mStagingManager.getStagedApexInfos(parentSession);
+        assertThat(result).containsExactly(fakeApexInfos[0].moduleName, fakeApexInfos[0],
+                fakeApexInfos[1].moduleName, fakeApexInfos[1]);
+
+        ArgumentCaptor<ApexSessionParams> argumentCaptor =
+                ArgumentCaptor.forClass(ApexSessionParams.class);
+        verify(mApexManager, times(1)).getStagedApexInfos(argumentCaptor.capture());
+        ApexSessionParams params = argumentCaptor.getValue();
+        assertThat(params.sessionId).isEqualTo(239);
+        assertThat(params.childSessionIds).asList().containsExactly(201, 202);
+    }
+
+    @Test
+    public void getStagedApexModuleNames_returnsStagedApexModules() throws Exception {
+        FakeStagedSession validSession1 = new FakeStagedSession(239);
+        validSession1.setIsApex(true);
+        validSession1.setSessionReady();
+        mStagingManager.createSession(validSession1);
+
+        FakeStagedSession childSession1 = new FakeStagedSession(123);
+        childSession1.setIsApex(true);
+        FakeStagedSession childSession2 = new FakeStagedSession(124);
+        childSession2.setIsApex(true);
+        FakeStagedSession nonApexChild = new FakeStagedSession(125);
+        FakeStagedSession parentSession = new FakeStagedSession(240,
+                Arrays.asList(childSession1, childSession2, nonApexChild));
+        parentSession.setSessionReady();
+        mStagingManager.createSession(parentSession);
+
+        // Make mApexManager return ApexInfo with same module name as the sessionId
+        // of the parameter that was passed into it
+        when(mApexManager.getStagedApexInfos(any())).thenAnswer(new Answer<ApexInfo[]>() {
+            @Override
+            public ApexInfo[] answer(InvocationOnMock invocation) throws Throwable {
+                Object[] args = invocation.getArguments();
+                ApexSessionParams params = (ApexSessionParams) args[0];
+                IntArray sessionsToProcess = new IntArray();
+                if (params.childSessionIds.length == 0) {
+                    sessionsToProcess.add(params.sessionId);
+                } else {
+                    sessionsToProcess.addAll(params.childSessionIds);
+                }
+                List<ApexInfo> result = new ArrayList<>();
+                for (int session : sessionsToProcess.toArray()) {
+                    ApexInfo info = new ApexInfo();
+                    info.moduleName = String.valueOf(session);
+                    result.add(info);
+                }
+                return result.toArray(new ApexInfo[0]);
+            }
+        });
+
+        List<String> result = mStagingManager.getStagedApexModuleNames();
+        assertThat(result).containsExactly("239", "123", "124");
+        verify(mApexManager, times(2)).getStagedApexInfos(any());
+    }
+
+    @Test
+    public void getStagedApexInfo() throws Exception {
+        FakeStagedSession validSession1 = new FakeStagedSession(239);
+        validSession1.setIsApex(true);
+        validSession1.setSessionReady();
+        mStagingManager.createSession(validSession1);
+        ApexInfo[] fakeApexInfos = getFakeApexInfo(Arrays.asList("module1"));
+        when(mApexManager.getStagedApexInfos(any())).thenReturn(fakeApexInfos);
+
+        // Verify null is returned if module name is not found
+        ApexInfo result = mStagingManager.getStagedApexInfo("not found");
+        assertThat(result).isNull();
+        verify(mApexManager, times(1)).getStagedApexInfos(any());
+        // Otherwise, the correct object is returned
+        result = mStagingManager.getStagedApexInfo("module1");
+        assertThat(result).isEqualTo(fakeApexInfos[0]);
+        verify(mApexManager, times(2)).getStagedApexInfos(any());
+    }
+
     private StagingManager.StagedSession createSession(int sessionId, String packageName,
             long committedMillis) {
         PackageInstaller.SessionParams params = new PackageInstaller.SessionParams(
@@ -582,10 +757,16 @@ public class StagingManagerTest {
         private String mPackageName;
         private boolean mIsAbandonded = false;
         private boolean mPreRebootVerificationStarted = false;
-        private final List<StagingManager.StagedSession> mChildSessions = new ArrayList<>();
+        private final List<StagingManager.StagedSession> mChildSessions;
 
         private FakeStagedSession(int sessionId) {
             mSessionId = sessionId;
+            mChildSessions = new ArrayList<>();
+        }
+
+        private FakeStagedSession(int sessionId, List<StagingManager.StagedSession> childSessions) {
+            mSessionId = sessionId;
+            mChildSessions = childSessions;
         }
 
         private void setParentSessionId(int parentSessionId) {