Merge "Add a utility class to track external callbacks"

6 files changed, 393 insertions, 108 deletions

diff --git a/service/java/com/android/server/wifi/WifiServiceImpl.java b/service/java/com/android/server/wifi/WifiServiceImpl.java
index a805f4a49..bf481edf7 100644
--- a/service/java/com/android/server/wifi/WifiServiceImpl.java
+++ b/service/java/com/android/server/wifi/WifiServiceImpl.java
@@ -100,6 +100,7 @@ import com.android.internal.telephony.PhoneConstants;
 import com.android.internal.telephony.TelephonyIntents;
 import com.android.internal.util.AsyncChannel;
 import com.android.server.wifi.hotspot2.PasspointProvider;
+import com.android.server.wifi.util.ExternalCallbackTracker;
 import com.android.server.wifi.util.GeneralUtil.Mutable;
 import com.android.server.wifi.util.WifiHandler;
 import com.android.server.wifi.util.WifiPermissionsUtil;
@@ -196,10 +197,7 @@ public class WifiServiceImpl extends IWifiManager.Stub {
     @GuardedBy("mLocalOnlyHotspotRequests")
     private final ConcurrentHashMap<String, Integer> mIfaceIpModes;
 
-    /* Limit on number of registered soft AP callbacks to track and prevent potential memory leak */
-    private static final int NUM_SOFT_AP_CALLBACKS_WARN_LIMIT = 10;
-    private static final int NUM_SOFT_AP_CALLBACKS_WTF_LIMIT = 20;
-    private final HashMap<Integer, ISoftApCallback> mRegisteredSoftApCallbacks;
+    private final ExternalCallbackTracker<ISoftApCallback> mRegisteredSoftApCallbacks;
 
     /**
      * One of:  {@link WifiManager#WIFI_AP_STATE_DISABLED},
@@ -457,7 +455,8 @@ public class WifiServiceImpl extends IWifiManager.Stub {
         mIfaceIpModes = new ConcurrentHashMap<>();
         mLocalOnlyHotspotRequests = new HashMap<>();
         enableVerboseLoggingInternal(getVerboseLoggingLevel());
-        mRegisteredSoftApCallbacks = new HashMap<>();
+        mRegisteredSoftApCallbacks =
+                new ExternalCallbackTracker<ISoftApCallback>(mClientModeImplHandler);
 
         mWifiInjector.getActiveModeWarden().registerSoftApCallback(new SoftApCallbackImpl());
         mPowerProfile = mWifiInjector.getPowerProfile();
@@ -1070,7 +1069,8 @@ public class WifiServiceImpl extends IWifiManager.Stub {
         public void onStateChanged(int state, int failureReason) {
             mSoftApState = state;
 
-            Iterator<ISoftApCallback> iterator = mRegisteredSoftApCallbacks.values().iterator();
+            Iterator<ISoftApCallback> iterator =
+                    mRegisteredSoftApCallbacks.getCallbacks().iterator();
             while (iterator.hasNext()) {
                 ISoftApCallback callback = iterator.next();
                 try {
@@ -1091,7 +1091,8 @@ public class WifiServiceImpl extends IWifiManager.Stub {
         public void onNumClientsChanged(int numClients) {
             mSoftApNumClients = numClients;
 
-            Iterator<ISoftApCallback> iterator = mRegisteredSoftApCallbacks.values().iterator();
+            Iterator<ISoftApCallback> iterator =
+                    mRegisteredSoftApCallbacks.getCallbacks().iterator();
             while (iterator.hasNext()) {
                 ISoftApCallback callback = iterator.next();
                 try {
@@ -1132,33 +1133,12 @@ public class WifiServiceImpl extends IWifiManager.Stub {
             mLog.info("registerSoftApCallback uid=%").c(Binder.getCallingUid()).flush();
         }
 
-        // register for binder death
-        IBinder.DeathRecipient dr = new IBinder.DeathRecipient() {
-            @Override
-            public void binderDied() {
-                binder.unlinkToDeath(this, 0);
-                mClientHandler.post(() -> {
-                    mRegisteredSoftApCallbacks.remove(callbackIdentifier);
-                });
-            }
-        };
-        try {
-            binder.linkToDeath(dr, 0);
-        } catch (RemoteException e) {
-            Log.e(TAG, "Error on linkToDeath - " + e);
-            return;
-        }
-
         // post operation to handler thread
         mClientHandler.post(() -> {
-            mRegisteredSoftApCallbacks.put(callbackIdentifier, callback);
-
-            if (mRegisteredSoftApCallbacks.size() > NUM_SOFT_AP_CALLBACKS_WTF_LIMIT) {
-                Log.wtf(TAG, "Too many soft AP callbacks: " + mRegisteredSoftApCallbacks.size());
-            } else if (mRegisteredSoftApCallbacks.size() > NUM_SOFT_AP_CALLBACKS_WARN_LIMIT) {
-                Log.w(TAG, "Too many soft AP callbacks: " + mRegisteredSoftApCallbacks.size());
+            if (!mRegisteredSoftApCallbacks.add(binder, callback, callbackIdentifier)) {
+                Log.e(TAG, "registerSoftApCallback: Failed to add callback");
+                return;
             }
-
             // Update the client about the current state immediately after registering the callback
             try {
                 callback.onStateChanged(mSoftApState, 0);
@@ -2871,26 +2851,9 @@ public class WifiServiceImpl extends IWifiManager.Stub {
         if (mVerboseLoggingEnabled) {
             mLog.info("registerTrafficStateCallback uid=%").c(Binder.getCallingUid()).flush();
         }
-
-        // register for binder death
-        IBinder.DeathRecipient dr = new IBinder.DeathRecipient() {
-            @Override
-            public void binderDied() {
-                binder.unlinkToDeath(this, 0);
-                mClientHandler.post(() -> {
-                    mTrafficPoller.removeCallback(callbackIdentifier);
-                });
-            }
-        };
-        try {
-            binder.linkToDeath(dr, 0);
-        } catch (RemoteException e) {
-            Log.e(TAG, "Error on linkToDeath - " + e);
-            return;
-        }
         // Post operation to handler thread
         mClientHandler.post(() -> {
-            mTrafficPoller.addCallback(callback, callbackIdentifier);
+            mTrafficPoller.addCallback(binder, callback, callbackIdentifier);
         });
     }
 
@@ -2908,7 +2871,6 @@ public class WifiServiceImpl extends IWifiManager.Stub {
         if (mVerboseLoggingEnabled) {
             mLog.info("unregisterTrafficStateCallback uid=%").c(Binder.getCallingUid()).flush();
         }
-
         // Post operation to handler thread
         mClientHandler.post(() -> {
             mTrafficPoller.removeCallback(callbackIdentifier);
diff --git a/service/java/com/android/server/wifi/WifiTrafficPoller.java b/service/java/com/android/server/wifi/WifiTrafficPoller.java
index d7f02db89..e484bf461 100644
--- a/service/java/com/android/server/wifi/WifiTrafficPoller.java
+++ b/service/java/com/android/server/wifi/WifiTrafficPoller.java
@@ -27,15 +27,17 @@ import android.net.NetworkInfo;
 import android.net.wifi.ITrafficStateCallback;
 import android.net.wifi.WifiManager;
 import android.os.Handler;
+import android.os.IBinder;
 import android.os.Looper;
 import android.os.Message;
 import android.os.RemoteException;
 import android.text.TextUtils;
 import android.util.Log;
 
+import com.android.server.wifi.util.ExternalCallbackTracker;
+
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
-import java.util.HashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
@@ -50,9 +52,6 @@ public class WifiTrafficPoller {
      * statistics
      */
     private static final int POLL_TRAFFIC_STATS_INTERVAL_MSECS = 1000;
-    /* Limit on number of registered soft AP callbacks to track and prevent potential memory leak */
-    private static final int NUM_CALLBACKS_WARN_LIMIT = 10;
-    private static final int NUM_CALLBACKS_WTF_LIMIT = 20;
 
     private static final int ENABLE_TRAFFIC_STATS_POLL  = 1;
     private static final int TRAFFIC_STATS_POLL         = 2;
@@ -64,7 +63,7 @@ public class WifiTrafficPoller {
     /* Tracks last reported data activity */
     private int mDataActivity;
 
-    private final HashMap<Integer, ITrafficStateCallback> mRegisteredCallbacks = new HashMap<>();
+    private final ExternalCallbackTracker<ITrafficStateCallback> mRegisteredCallbacks;
     // err on the side of updating at boot since screen on broadcast may be missed
     // the first time
     private AtomicBoolean mScreenOn = new AtomicBoolean(true);
@@ -75,9 +74,10 @@ public class WifiTrafficPoller {
     private boolean mVerboseLoggingEnabled = false;
 
     WifiTrafficPoller(@NonNull Context context, @NonNull Looper looper,
-            @NonNull WifiNative wifiNative) {
+                      @NonNull WifiNative wifiNative) {
         mTrafficHandler = new TrafficHandler(looper);
         mWifiNative = wifiNative;
+        mRegisteredCallbacks = new ExternalCallbackTracker<ITrafficStateCallback>(mTrafficHandler);
 
         IntentFilter filter = new IntentFilter();
         filter.addAction(WifiManager.NETWORK_STATE_CHANGED_ACTION);
@@ -108,25 +108,25 @@ public class WifiTrafficPoller {
     /**
      * Add a new callback to the traffic poller.
      */
-    public void addCallback(ITrafficStateCallback callback, int callbackIdentifier) {
-        mRegisteredCallbacks.put(callbackIdentifier, callback);
-        if (mVerboseLoggingEnabled) {
-            Log.d(TAG, "Adding callback. Num callbacks: " + mRegisteredCallbacks.size());
+    public void addCallback(IBinder binder, ITrafficStateCallback callback,
+                            int callbackIdentifier) {
+        if (!mRegisteredCallbacks.add(binder, callback, callbackIdentifier)) {
+            Log.e(TAG, "Failed to add callback");
+            return;
         }
-        if (mRegisteredCallbacks.size() > NUM_CALLBACKS_WTF_LIMIT) {
-            Log.wtf(TAG, "Too many traffic poller callbacks: " + mRegisteredCallbacks.size());
-        } else if (mRegisteredCallbacks.size() > NUM_CALLBACKS_WARN_LIMIT) {
-            Log.w(TAG, "Too many traffic poller callbacks: " + mRegisteredCallbacks.size());
+        if (mVerboseLoggingEnabled) {
+            Log.v(TAG, "Adding callback. Num callbacks: " + mRegisteredCallbacks.getNumCallbacks());
         }
     }
 
     /**
-     * Remove an existing callback to the traffic poller.
+     * Remove an existing callback from the traffic poller.
      */
     public void removeCallback(int callbackIdentifier) {
         mRegisteredCallbacks.remove(callbackIdentifier);
         if (mVerboseLoggingEnabled) {
-            Log.d(TAG, "Removing callback. Num callbacks: " + mRegisteredCallbacks.size());
+            Log.v(TAG, "Removing callback. Num callbacks: "
+                    + mRegisteredCallbacks.getNumCallbacks());
         }
     }
 
@@ -162,11 +162,11 @@ public class WifiTrafficPoller {
                     }
                     break;
                 case TRAFFIC_STATS_POLL:
-                    if (DBG) {
+                    if (mVerboseLoggingEnabled) {
                         Log.d(TAG, "TRAFFIC_STATS_POLL "
                                 + mEnableTrafficStatsPoll + " Token "
                                 + Integer.toString(mTrafficStatsPollToken)
-                                + " num clients " + mRegisteredCallbacks.size());
+                                + " num clients " + mRegisteredCallbacks.getNumCallbacks());
                     }
                     if (msg.arg1 == mTrafficStatsPollToken) {
                         ifaceName = mWifiNative.getClientInterfaceName();
@@ -226,7 +226,7 @@ public class WifiTrafficPoller {
                     Log.e(TAG, "notifying of data activity "
                             + Integer.toString(mDataActivity));
                 }
-                for (ITrafficStateCallback callback : mRegisteredCallbacks.values()) {
+                for (ITrafficStateCallback callback : mRegisteredCallbacks.getCallbacks()) {
                     try {
                         callback.onStateChanged(mDataActivity);
                     } catch (RemoteException e) {
@@ -238,13 +238,16 @@ public class WifiTrafficPoller {
         }
     }
 
-    void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
+    /**
+     * Dump method for traffic poller.
+     */
+    public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
         pw.println("mEnableTrafficStatsPoll " + mEnableTrafficStatsPoll);
         pw.println("mTrafficStatsPollToken " + mTrafficStatsPollToken);
         pw.println("mTxPkts " + mTxPkts);
         pw.println("mRxPkts " + mRxPkts);
         pw.println("mDataActivity " + mDataActivity);
-        pw.println("mRegisteredCallbacks " + mRegisteredCallbacks);
+        pw.println("mRegisteredCallbacks " + mRegisteredCallbacks.getNumCallbacks());
     }
 
 }
diff --git a/service/java/com/android/server/wifi/util/ExternalCallbackTracker.java b/service/java/com/android/server/wifi/util/ExternalCallbackTracker.java
new file mode 100644
index 000000000..22132c436
--- /dev/null
+++ b/service/java/com/android/server/wifi/util/ExternalCallbackTracker.java
@@ -0,0 +1,174 @@
+/*
+ * 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 android.annotation.NonNull;
+import android.os.Handler;
+import android.os.IBinder;
+import android.os.RemoteException;
+import android.util.Log;
+
+import com.android.internal.util.Preconditions;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Holds a list of external app-provided binder callback objects and tracks the death
+ * of the callback object.
+ * @param <T> Callback object type.
+ */
+public class ExternalCallbackTracker<T> {
+    private static final String TAG = "WifiExternalCallbackTracker";
+
+    /* Limit on number of registered callbacks to track and prevent potential memory leak */
+    private static final int NUM_CALLBACKS_WARN_LIMIT = 10;
+    private static final int NUM_CALLBACKS_WTF_LIMIT = 20;
+
+    /**
+     * Container for storing info about each external callback and tracks it's death.
+     */
+    private static class ExternalCallbackHolder<T> implements IBinder.DeathRecipient {
+        private final IBinder mBinder;
+        private final T mCallbackObject;
+        private final DeathCallback mDeathCallback;
+
+        /**
+         * Callback to be invoked on death of the app hosting the binder.
+         */
+        public interface DeathCallback {
+            /**
+             * Called when the corresponding app has died.
+             */
+            void onDeath();
+        }
+
+        private ExternalCallbackHolder(@NonNull IBinder binder, @NonNull T callbackObject,
+                                       @NonNull DeathCallback deathCallback) {
+            mBinder = Preconditions.checkNotNull(binder);
+            mCallbackObject = Preconditions.checkNotNull(callbackObject);
+            mDeathCallback = Preconditions.checkNotNull(deathCallback);
+        }
+
+        /**
+         * Static method to create a new {@link ExternalCallbackHolder} object and register for
+         * death notification of the associated binder.
+         * @return an instance of {@link ExternalCallbackHolder} if there are no failures, otherwise
+         * null.
+         */
+        public static <T> ExternalCallbackHolder createAndLinkToDeath(
+                @NonNull IBinder binder, @NonNull T callbackObject,
+                @NonNull DeathCallback deathCallback) {
+            ExternalCallbackHolder<T> externalCallback =
+                    new ExternalCallbackHolder<T>(binder, callbackObject, deathCallback);
+            try {
+                binder.linkToDeath(externalCallback, 0);
+            } catch (RemoteException e) {
+                Log.e(TAG, "Error on linkToDeath - " + e);
+                return null;
+            }
+            return externalCallback;
+        }
+
+        /**
+         * Unlinks this object from binder death.
+         */
+        public void reset() {
+            mBinder.unlinkToDeath(this, 0);
+        }
+
+        /**
+         * Retrieve the callback object.
+         */
+        public T getCallback() {
+            return mCallbackObject;
+        }
+
+        /**
+         * App hosting the binder has died.
+         */
+        @Override
+        public void binderDied() {
+            mDeathCallback.onDeath();
+            Log.d(TAG, "Binder died " + mBinder);
+        }
+    }
+
+    private final Map<Integer, ExternalCallbackHolder<T>> mCallbacks;
+    private final Handler mHandler;
+
+    public ExternalCallbackTracker(Handler handler) {
+        mHandler = handler;
+        mCallbacks = new HashMap<>();
+    }
+
+    /**
+     * Add a callback object to tracker.
+     * @return true on success, false on failure.
+     */
+    public boolean add(@NonNull IBinder binder, @NonNull T callbackObject, int callbackIdentifier) {
+        ExternalCallbackHolder<T> externalCallback = ExternalCallbackHolder.createAndLinkToDeath(
+                binder, callbackObject, () -> {
+                    mHandler.post(() -> {
+                        Log.d(TAG, "Remove external callback on death " + callbackIdentifier);
+                        remove(callbackIdentifier);
+                    });
+                });
+        if (externalCallback == null) return false;
+        mCallbacks.put(callbackIdentifier, externalCallback);
+        if (mCallbacks.size() > NUM_CALLBACKS_WTF_LIMIT) {
+            Log.wtf(TAG, "Too many callbacks: " + mCallbacks.size());
+        } else if (mCallbacks.size() > NUM_CALLBACKS_WARN_LIMIT) {
+            Log.w(TAG, "Too many callbacks: " + mCallbacks.size());
+        }
+        return true;
+    }
+
+    /**
+     * Remove a callback object to tracker.
+     * @return true on success, false on failure.
+     */
+    public boolean remove(int callbackIdentifier) {
+        ExternalCallbackHolder<T> externalCallback = mCallbacks.remove(callbackIdentifier);
+        if (externalCallback == null) {
+            Log.w(TAG, "Unknown external callback " + callbackIdentifier);
+            return false;
+        }
+        externalCallback.reset();
+        return true;
+    }
+
+    /**
+     * Retrieve all the callback objects in the tracker.
+     */
+    public List<T> getCallbacks() {
+        List<T> callbacks = new ArrayList<>();
+        for (ExternalCallbackHolder<T> externalCallback : mCallbacks.values()) {
+            callbacks.add(externalCallback.getCallback());
+        }
+        return callbacks;
+    }
+
+    /**
+     * Retrieve the number of callback objects in the tracker.
+     */
+    public int getNumCallbacks() {
+        return mCallbacks.size();
+    }
+}
diff --git a/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java b/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java
index 4eedf8384..478594dc5 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java
@@ -1447,6 +1447,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.
      */
@@ -2924,7 +2939,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);
     }
 
     /**
@@ -2936,32 +2951,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());
+    }
+}