Use NetworkStack.getService in NetworkStackClient

This avoids returning IBinder in Context.getSystemService, having a
dedicated getter for the NetworkStack AIDL interface instead.

Test: atest FrameworksNetTests NetworkStackTests
Bug: 151243982
Change-Id: Ibeaeff5890631e14581a084863d9aeb3178c5d5c

2 files changed, 21 insertions, 17 deletions

diff --git a/common/networkstackclient/src/android/net/networkstack/ModuleNetworkStackClient.java b/common/networkstackclient/src/android/net/networkstack/ModuleNetworkStackClient.java
index cfbb760..cbbae37 100644
--- a/common/networkstackclient/src/android/net/networkstack/ModuleNetworkStackClient.java
+++ b/common/networkstackclient/src/android/net/networkstack/ModuleNetworkStackClient.java
@@ -16,12 +16,12 @@
 
 package android.net.networkstack;
 
-import static android.content.Context.NETWORK_STACK_SERVICE;
 import static android.os.Build.VERSION.SDK_INT;
 
 import android.annotation.NonNull;
 import android.content.Context;
 import android.net.INetworkStackConnector;
+import android.net.NetworkStack;
 import android.os.Build;
 import android.os.IBinder;
 import android.util.Log;
@@ -47,14 +47,14 @@ public class ModuleNetworkStackClient extends NetworkStackClientBase {
         // TODO(b/149676685): change this check to "< R" once R is defined
         if (SDK_INT < Build.VERSION_CODES.Q
                 || (SDK_INT == Build.VERSION_CODES.Q && "REL".equals(Build.VERSION.CODENAME))) {
-            // NETWORK_STACK_SERVICE is not available through getSystemService before R
+            // The NetworkStack connector is not available through NetworkStack before R
             throw new UnsupportedOperationException(
                     "ModuleNetworkStackClient is not supported on API " + SDK_INT);
         }
 
         if (sInstance == null) {
             sInstance = new ModuleNetworkStackClient();
-            sInstance.startPolling(packageContext);
+            sInstance.startPolling();
         }
         return sInstance;
     }
@@ -64,10 +64,10 @@ public class ModuleNetworkStackClient extends NetworkStackClientBase {
         sInstance = null;
     }
 
-    private void startPolling(Context context) {
+    private void startPolling() {
         // If the service is already registered (as it will be most of the time), do not poll and
         // fulfill requests immediately.
-        final IBinder nss = (IBinder) context.getSystemService(NETWORK_STACK_SERVICE);
+        final IBinder nss = NetworkStack.getService();
         if (nss != null) {
             // Calling onNetworkStackConnected here means that pending oneway Binder calls to the
             // NetworkStack get sent from the current thread and not a worker thread; this is fine
@@ -75,21 +75,18 @@ public class ModuleNetworkStackClient extends NetworkStackClientBase {
             onNetworkStackConnected(INetworkStackConnector.Stub.asInterface(nss));
             return;
         }
-        new Thread(new PollingRunner(context)).start();
+        new Thread(new PollingRunner()).start();
     }
 
     private class PollingRunner implements Runnable {
-        private final Context mContext;
 
-        private PollingRunner(Context context) {
-            mContext = context;
-        }
+        private PollingRunner() {}
 
         @Override
         public void run() {
             // Block until the NetworkStack connector is registered in ServiceManager.
             IBinder nss;
-            while ((nss = (IBinder) mContext.getSystemService(NETWORK_STACK_SERVICE)) == null) {
+            while ((nss = NetworkStack.getService()) == null) {
                 try {
                     Thread.sleep(200);
                 } catch (InterruptedException e) {
diff --git a/tests/unit/src/android/net/networkstack/ModuleNetworkStackClientTest.kt b/tests/unit/src/android/net/networkstack/ModuleNetworkStackClientTest.kt
index 661f9aa..1024af7 100644
--- a/tests/unit/src/android/net/networkstack/ModuleNetworkStackClientTest.kt
+++ b/tests/unit/src/android/net/networkstack/ModuleNetworkStackClientTest.kt
@@ -21,6 +21,7 @@ import android.net.IIpMemoryStoreCallbacks
 import android.net.INetworkMonitorCallbacks
 import android.net.INetworkStackConnector
 import android.net.Network
+import android.net.NetworkStack
 import android.net.dhcp.DhcpServingParamsParcel
 import android.net.dhcp.IDhcpServerCallbacks
 import android.net.ip.IIpClientCallbacks
@@ -28,6 +29,8 @@ import android.os.Build
 import android.os.IBinder
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn
+import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
 import com.android.networkstack.apishim.ShimUtils
 import org.junit.After
 import org.junit.Assume.assumeTrue
@@ -36,11 +39,11 @@ import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentMatchers.any
 import org.mockito.Mock
-import org.mockito.Mockito.doReturn
 import org.mockito.Mockito.never
 import org.mockito.Mockito.timeout
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
+import org.mockito.MockitoSession
 
 @RunWith(AndroidJUnit4::class)
 @SmallTest
@@ -49,6 +52,8 @@ class ModuleNetworkStackClientTest {
     private val TEST_NETWORK = Network(43)
     private val TEST_TIMEOUT_MS = 500L
 
+    private var mSession: MockitoSession? = null
+
     @Mock
     private lateinit var mContext: Context
     @Mock
@@ -68,6 +73,7 @@ class ModuleNetworkStackClientTest {
     fun setUp() {
         // ModuleNetworkStackClient is only available after Q
         assumeTrue(ShimUtils.isReleaseOrDevelopmentApiAbove(Build.VERSION_CODES.Q))
+        mSession = mockitoSession().spyStatic(NetworkStack::class.java).startMocking()
         MockitoAnnotations.initMocks(this)
         doReturn(mConnector).`when`(mConnectorBinder).queryLocalInterface(
                 INetworkStackConnector::class.qualifiedName!!)
@@ -75,12 +81,13 @@ class ModuleNetworkStackClientTest {
 
     @After
     fun tearDown() {
+        mSession?.finishMocking()
         ModuleNetworkStackClient.resetInstanceForTest()
     }
 
     @Test
     fun testIpClientServiceAvailableImmediately() {
-        doReturn(mConnectorBinder).`when`(mContext).getSystemService(Context.NETWORK_STACK_SERVICE)
+        doReturn(mConnectorBinder).`when` { NetworkStack.getService() }
         ModuleNetworkStackClient.getInstance(mContext).makeIpClient(TEST_IFNAME, mIpClientCb)
         verify(mConnector).makeIpClient(TEST_IFNAME, mIpClientCb)
     }
@@ -91,14 +98,14 @@ class ModuleNetworkStackClientTest {
 
         Thread.sleep(TEST_TIMEOUT_MS)
         verify(mConnector, never()).makeIpClient(any(), any())
-        doReturn(mConnectorBinder).`when`(mContext).getSystemService(Context.NETWORK_STACK_SERVICE)
+        doReturn(mConnectorBinder).`when` { NetworkStack.getService() }
         // Use a longer timeout as polling can cause larger delays
         verify(mConnector, timeout(TEST_TIMEOUT_MS * 4)).makeIpClient(TEST_IFNAME, mIpClientCb)
     }
 
     @Test
     fun testDhcpServerAvailableImmediately() {
-        doReturn(mConnectorBinder).`when`(mContext).getSystemService(Context.NETWORK_STACK_SERVICE)
+        doReturn(mConnectorBinder).`when` { NetworkStack.getService() }
         val testParams = DhcpServingParamsParcel()
         ModuleNetworkStackClient.getInstance(mContext).makeDhcpServer(TEST_IFNAME, testParams,
                 mDhcpServerCb)
@@ -107,7 +114,7 @@ class ModuleNetworkStackClientTest {
 
     @Test
     fun testNetworkMonitorAvailableImmediately() {
-        doReturn(mConnectorBinder).`when`(mContext).getSystemService(Context.NETWORK_STACK_SERVICE)
+        doReturn(mConnectorBinder).`when` { NetworkStack.getService() }
         val testName = "NetworkMonitorName"
         ModuleNetworkStackClient.getInstance(mContext).makeNetworkMonitor(TEST_NETWORK, testName,
                 mNetworkMonitorCb)
@@ -116,7 +123,7 @@ class ModuleNetworkStackClientTest {
 
     @Test
     fun testIpMemoryStoreAvailableImmediately() {
-        doReturn(mConnectorBinder).`when`(mContext).getSystemService(Context.NETWORK_STACK_SERVICE)
+        doReturn(mConnectorBinder).`when` { NetworkStack.getService() }
         ModuleNetworkStackClient.getInstance(mContext).fetchIpMemoryStore(mIpMemoryStoreCb)
         verify(mConnector).fetchIpMemoryStore(mIpMemoryStoreCb)
     }