Add a utility class to track external callbacks

Add a new class: ExternalCallbackTracker to keep track of external
binder callback objects. This class will:
a) Hold a list of callback instances.
b) Register for death notification on the associated binder for each
callback.
c) Will automatically remove the callback from the list on binder death.

Moved ISoftApCallback & ITrafficStateCallback object registrations to
use this new object. This fixes a leak (not unlinking death recipient on
remove) with the existing callback management.

Bug: 27074039
Test: Unit tests
Test: Verified manually that softap & traffic state change callbacks are
still working.

Change-Id: Ifee9504e58befc2f8e3b575c8dc01d484069387d

3 files changed, 183 insertions, 37 deletions

diff --git a/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java b/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java
index f4a97c2b0..761518174 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java
@@ -1451,6 +1451,21 @@ public class WifiServiceImplTest {
     }
 
     /**
+     * Verifies that we handle softap callback registration failure if we encounter an exception
+     * while linking to death.
+     */
+    @Test
+    public void registerSoftApCallbackFailureOnLinkToDeath() throws Exception {
+        doThrow(new RemoteException())
+                .when(mAppBinder).linkToDeath(any(IBinder.DeathRecipient.class), anyInt());
+        mWifiServiceImpl.registerSoftApCallback(mAppBinder, mClientSoftApCallback, 1);
+        mLooper.dispatchAll();
+        verify(mClientSoftApCallback, never()).onStateChanged(WIFI_AP_STATE_DISABLED, 0);
+        verify(mClientSoftApCallback, never()).onNumClientsChanged(0);
+    }
+
+
+    /**
      * Registers a soft AP callback, then verifies that the current soft AP state and num clients
      * are sent to caller immediately after callback is registered.
      */
@@ -2925,7 +2940,7 @@ public class WifiServiceImplTest {
                 mAppBinder, mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
         mLooper.dispatchAll();
         verify(mWifiTrafficPoller).addCallback(
-                mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+                mAppBinder, mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
     }
 
     /**
@@ -2937,32 +2952,4 @@ public class WifiServiceImplTest {
         mLooper.dispatchAll();
         verify(mWifiTrafficPoller).removeCallback(0);
     }
-
-    /**
-     * Verify that wifi service registers for callers BinderDeath event
-     */
-    @Test
-    public void registersForBinderDeathOnRegisterTrafficStateCallback() throws Exception {
-        mWifiServiceImpl.registerTrafficStateCallback(
-                mAppBinder, mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
-        mLooper.dispatchAll();
-        verify(mAppBinder).linkToDeath(any(IBinder.DeathRecipient.class), anyInt());
-    }
-
-    /**
-     * Verify that we remove the traffic state callback on receiving BinderDied event.
-     */
-    @Test
-    public void unregistersTrafficStateCallbackOnBinderDied() throws Exception {
-        ArgumentCaptor<IBinder.DeathRecipient> drCaptor =
-                ArgumentCaptor.forClass(IBinder.DeathRecipient.class);
-        mWifiServiceImpl.registerTrafficStateCallback(
-                mAppBinder, mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
-        verify(mAppBinder).linkToDeath(drCaptor.capture(), anyInt());
-
-        drCaptor.getValue().binderDied();
-        mLooper.dispatchAll();
-        verify(mAppBinder).unlinkToDeath(drCaptor.getValue(), 0);
-        verify(mWifiTrafficPoller).removeCallback(TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
-    }
 }
diff --git a/tests/wifitests/src/com/android/server/wifi/WifiTrafficPollerTest.java b/tests/wifitests/src/com/android/server/wifi/WifiTrafficPollerTest.java
index b00c4f74c..4115aabcd 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiTrafficPollerTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiTrafficPollerTest.java
@@ -19,6 +19,7 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -29,11 +30,13 @@ import android.content.Intent;
 import android.net.NetworkInfo;
 import android.net.wifi.ITrafficStateCallback;
 import android.net.wifi.WifiManager;
-import android.os.Message;
+import android.os.IBinder;
 import android.os.RemoteException;
 import android.os.test.TestLooper;
 import android.support.test.filters.SmallTest;
 
+import com.android.server.wifi.util.ExternalCallbackTracker;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
@@ -57,14 +60,15 @@ public class WifiTrafficPollerTest {
     private final static long RX_PACKET_COUNT = 50;
     private static final int TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER = 14;
 
-    final ArgumentCaptor<Message> mMessageCaptor = ArgumentCaptor.forClass(Message.class);
     final ArgumentCaptor<BroadcastReceiver> mBroadcastReceiverCaptor =
             ArgumentCaptor.forClass(BroadcastReceiver.class);
 
     @Mock Context mContext;
     @Mock WifiNative mWifiNative;
     @Mock NetworkInfo mNetworkInfo;
+    @Mock IBinder mAppBinder;
     @Mock ITrafficStateCallback mTrafficStateCallback;
+    @Mock ExternalCallbackTracker<ITrafficStateCallback> mCallbackTracker;
 
     /**
      * Called before each test
@@ -116,7 +120,7 @@ public class WifiTrafficPollerTest {
     public void testNotStartTrafficStatsPollingWithDisconnected() throws RemoteException {
         // Register Client to verify that Tx/RX packet message is properly received.
         mWifiTrafficPoller.addCallback(
-                mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+                mAppBinder, mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
         triggerForUpdatedInformationOfData(Intent.ACTION_SCREEN_ON,
                 NetworkInfo.DetailedState.DISCONNECTED);
 
@@ -132,7 +136,7 @@ public class WifiTrafficPollerTest {
     public void testStartTrafficStatsPollingWithScreenOn() throws RemoteException {
         // Register Client to verify that Tx/RX packet message is properly received.
         mWifiTrafficPoller.addCallback(
-                mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+                mAppBinder, mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
         triggerForUpdatedInformationOfData(Intent.ACTION_SCREEN_ON,
                 NetworkInfo.DetailedState.CONNECTED);
 
@@ -148,7 +152,7 @@ public class WifiTrafficPollerTest {
     public void testNotStartTrafficStatsPollingWithScreenOff() throws RemoteException {
         // Register Client to verify that Tx/RX packet message is properly received.
         mWifiTrafficPoller.addCallback(
-                mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+                mAppBinder, mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
         triggerForUpdatedInformationOfData(Intent.ACTION_SCREEN_OFF,
                 NetworkInfo.DetailedState.CONNECTED);
 
@@ -166,9 +170,9 @@ public class WifiTrafficPollerTest {
     public void testRemoveClient() throws RemoteException {
         // Register Client to verify that Tx/RX packet message is properly received.
         mWifiTrafficPoller.addCallback(
-                mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+                mAppBinder, mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
         mWifiTrafficPoller.removeCallback(TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
-        mLooper.dispatchAll();
+        verify(mAppBinder).unlinkToDeath(any(), anyInt());
 
         triggerForUpdatedInformationOfData(Intent.ACTION_SCREEN_ON,
                 NetworkInfo.DetailedState.CONNECTED);
@@ -184,7 +188,7 @@ public class WifiTrafficPollerTest {
     public void testRemoveClientWithWrongIdentifier() throws RemoteException {
         // Register Client to verify that Tx/RX packet message is properly received.
         mWifiTrafficPoller.addCallback(
-                mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+                mAppBinder, mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
         mWifiTrafficPoller.removeCallback(TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER + 5);
         mLooper.dispatchAll();
 
@@ -195,4 +199,34 @@ public class WifiTrafficPollerTest {
         verify(mTrafficStateCallback).onStateChanged(
                 WifiManager.TrafficStateCallback.DATA_ACTIVITY_INOUT);
     }
+
+    /**
+     *
+     * Verify that traffic poller registers for death notification on adding client.
+     */
+    @Test
+    public void registersForBinderDeathOnAddClient() throws Exception {
+        mWifiTrafficPoller.addCallback(
+                mAppBinder, mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+        verify(mAppBinder).linkToDeath(any(IBinder.DeathRecipient.class), anyInt());
+    }
+
+    /**
+     *
+     * Verify that traffic poller registers for death notification on adding client.
+     */
+    @Test
+    public void addCallbackFailureOnLinkToDeath() throws Exception {
+        doThrow(new RemoteException())
+                .when(mAppBinder).linkToDeath(any(IBinder.DeathRecipient.class), anyInt());
+        mWifiTrafficPoller.addCallback(
+                mAppBinder, mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+        verify(mAppBinder).linkToDeath(any(IBinder.DeathRecipient.class), anyInt());
+
+        triggerForUpdatedInformationOfData(Intent.ACTION_SCREEN_ON,
+                NetworkInfo.DetailedState.CONNECTED);
+
+        // Client should not get any message callback add failed.
+        verify(mTrafficStateCallback, never()).onStateChanged(anyInt());
+    }
 }
diff --git a/tests/wifitests/src/com/android/server/wifi/util/ExternalCallbackTrackerTest.java b/tests/wifitests/src/com/android/server/wifi/util/ExternalCallbackTrackerTest.java
new file mode 100644
index 000000000..94c1aee38
--- /dev/null
+++ b/tests/wifitests/src/com/android/server/wifi/util/ExternalCallbackTrackerTest.java
@@ -0,0 +1,125 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.wifi.util;
+
+import static org.junit.Assert.*;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.verify;
+
+import android.net.wifi.ISoftApCallback;
+import android.os.Handler;
+import android.os.IBinder;
+import android.os.RemoteException;
+import android.os.test.TestLooper;
+import android.test.suitebuilder.annotation.SmallTest;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Unit tests for {@link com.android.server.wifi.util.ExternalCallbackTracker}.
+ */
+@SmallTest
+public class ExternalCallbackTrackerTest {
+    private static final int TEST_CALLBACK_IDENTIFIER = 56;
+    @Mock Handler mHandler;
+    @Mock ISoftApCallback mCallback;
+    @Mock IBinder mBinder;
+    private TestLooper mTestLooper;
+
+    private ExternalCallbackTracker<ISoftApCallback> mExternalCallbackTracker;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mTestLooper = new TestLooper();
+        mHandler = new Handler(mTestLooper.getLooper());
+        mExternalCallbackTracker = new ExternalCallbackTracker<ISoftApCallback>(mHandler);
+    }
+
+    /**
+     * Test adding a callback.
+     */
+    @Test
+    public void testAddCallback() throws Exception {
+        assertTrue(mExternalCallbackTracker.add(mBinder, mCallback, TEST_CALLBACK_IDENTIFIER));
+        assertEquals(1, mExternalCallbackTracker.getNumCallbacks());
+        assertEquals(mCallback, mExternalCallbackTracker.getCallbacks().get(0));
+        verify(mBinder).linkToDeath(any(), anyInt());
+    }
+
+    /**
+     * Test that adding a callback returns failure when binder death linking fails.
+     */
+    @Test
+    public void testAddCallbackFailureOnLinkToDeath() throws Exception {
+        doThrow(new RemoteException()).when(mBinder).linkToDeath(any(), anyInt());
+        assertFalse(mExternalCallbackTracker.add(mBinder, mCallback, TEST_CALLBACK_IDENTIFIER));
+        assertEquals(0, mExternalCallbackTracker.getNumCallbacks());
+        assertTrue(mExternalCallbackTracker.getCallbacks().isEmpty());
+    }
+
+    /**
+     * Test removing a callback.
+     */
+    @Test
+    public void testRemoveCallback() throws Exception {
+        testAddCallback();
+
+        assertTrue(mExternalCallbackTracker.remove(TEST_CALLBACK_IDENTIFIER));
+        assertEquals(0, mExternalCallbackTracker.getNumCallbacks());
+        assertTrue(mExternalCallbackTracker.getCallbacks().isEmpty());
+        verify(mBinder).unlinkToDeath(any(), anyInt());
+    }
+
+    /**
+     * Test removing a callback returns failure when the identifier provided doesn't match the one
+     * used to add the callback.
+     */
+    @Test
+    public void testRemoveCallbackFailureOnWrongIdentifier() throws Exception {
+        testAddCallback();
+
+        assertFalse(mExternalCallbackTracker.remove(TEST_CALLBACK_IDENTIFIER + 5));
+        assertEquals(1, mExternalCallbackTracker.getNumCallbacks());
+        assertEquals(mCallback, mExternalCallbackTracker.getCallbacks().get(0));
+    }
+
+    /**
+     * Test that the callback is automatically removed when the associated binder object is dead.
+     */
+    @Test
+    public void testCallbackRemovalOnDeath() throws Exception {
+        assertTrue(mExternalCallbackTracker.add(mBinder, mCallback, TEST_CALLBACK_IDENTIFIER));
+
+        // Trigger the death.
+        ArgumentCaptor<IBinder.DeathRecipient> deathCaptor =
+                ArgumentCaptor.forClass(IBinder.DeathRecipient.class);
+        verify(mBinder).linkToDeath(deathCaptor.capture(), anyInt());
+        deathCaptor.getValue().binderDied();
+        mTestLooper.dispatchAll();
+
+        assertEquals(0, mExternalCallbackTracker.getNumCallbacks());
+        assertTrue(mExternalCallbackTracker.getCallbacks().isEmpty());
+        verify(mBinder).unlinkToDeath(any(), anyInt());
+    }
+}