WifiServiceImpl: Add callback for traffic poller

The current mechanism for traffic poller directly exposes the wifi
service's internal message handler to privileged apps. Move away from
this mechanism and expose an API to let apps register for callbacks from
traffic poller. This API uses binder IPC for registration and
invocation of callbacks. Clients are automatically removed on binder
death.

Note: This is based on the existing SoftapCallback registration
mechanism.

Bug: 27074039
Test: Unit tests
Test: Verified the data indicators on Sysui
Change-Id: I6da6027893f58ddd7d3061cc7c1ccdeecb081eac

4 files changed, 260 insertions, 87 deletions

diff --git a/service/java/com/android/server/wifi/WifiServiceImpl.java b/service/java/com/android/server/wifi/WifiServiceImpl.java
index 30fd29514..a805f4a49 100644
--- a/service/java/com/android/server/wifi/WifiServiceImpl.java
+++ b/service/java/com/android/server/wifi/WifiServiceImpl.java
@@ -58,6 +58,7 @@ import android.net.NetworkUtils;
 import android.net.Uri;
 import android.net.ip.IpClient;
 import android.net.wifi.ISoftApCallback;
+import android.net.wifi.ITrafficStateCallback;
 import android.net.wifi.IWifiManager;
 import android.net.wifi.ScanResult;
 import android.net.wifi.WifiActivityEnergyInfo;
@@ -250,26 +251,6 @@ public class WifiServiceImpl extends IWifiManager.Stub {
         public void handleMessage(Message msg) {
             super.handleMessage(msg);
             switch (msg.what) {
-                case AsyncChannel.CMD_CHANNEL_HALF_CONNECTED: {
-                    if (msg.arg1 == AsyncChannel.STATUS_SUCCESSFUL) {
-                        Slog.d(TAG, "New client listening to asynchronous messages");
-                        // We track the clients by the Messenger
-                        // since it is expected to be always available
-                        mTrafficPoller.addClient(msg.replyTo);
-                    } else {
-                        Slog.e(TAG, "Client connection failure, error=" + msg.arg1);
-                    }
-                    break;
-                }
-                case AsyncChannel.CMD_CHANNEL_DISCONNECTED: {
-                    if (msg.arg1 == AsyncChannel.STATUS_SEND_UNSUCCESSFUL) {
-                        Slog.w(TAG, "Send failed, client connection lost");
-                    } else {
-                        Slog.w(TAG, "Client connection lost with reason: " + msg.arg1);
-                    }
-                    mTrafficPoller.removeClient(msg.replyTo);
-                    break;
-                }
                 case AsyncChannel.CMD_CHANNEL_FULL_CONNECTION: {
                     AsyncChannel ac = mFrameworkFacade.makeWifiAsyncChannel(TAG);
                     ac.connect(mContext, this, msg.replyTo);
@@ -2862,4 +2843,75 @@ public class WifiServiceImpl extends IWifiManager.Stub {
                     .c(provider.toString()).flush();
         }
     }
+
+    /**
+     * see {@link android.net.wifi.WifiManager#registerTrafficStateCallback(
+     * TrafficStateCallback, Handler)}
+     *
+     * @param binder IBinder instance to allow cleanup if the app dies
+     * @param callback Traffic State callback to register
+     * @param callbackIdentifier Unique ID of the registering callback. This ID will be used to
+     *        unregister the callback. See {@link unregisterTrafficStateCallback(int)}
+     *
+     * @throws SecurityException if the caller does not have permission to register a callback
+     * @throws RemoteException if remote exception happens
+     * @throws IllegalArgumentException if the arguments are null or invalid
+     */
+    @Override
+    public void registerTrafficStateCallback(IBinder binder, ITrafficStateCallback callback,
+                                             int callbackIdentifier) {
+        // verify arguments
+        if (binder == null) {
+            throw new IllegalArgumentException("Binder must not be null");
+        }
+        if (callback == null) {
+            throw new IllegalArgumentException("Callback must not be null");
+        }
+        enforceNetworkSettingsPermission();
+        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);
+        });
+    }
+
+    /**
+     * see {@link android.net.wifi.WifiManager#unregisterTrafficStateCallback(
+     * WifiManager.TrafficStateCallback)}
+     *
+     * @param callbackIdentifier Unique ID of the callback to be unregistered.
+     *
+     * @throws SecurityException if the caller does not have permission to register a callback
+     */
+    @Override
+    public void unregisterTrafficStateCallback(int callbackIdentifier) {
+        enforceNetworkSettingsPermission();
+        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 bcb625ed6..d7f02db89 100644
--- a/service/java/com/android/server/wifi/WifiTrafficPoller.java
+++ b/service/java/com/android/server/wifi/WifiTrafficPoller.java
@@ -24,19 +24,18 @@ import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
 import android.net.NetworkInfo;
+import android.net.wifi.ITrafficStateCallback;
 import android.net.wifi.WifiManager;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
-import android.os.Messenger;
 import android.os.RemoteException;
 import android.text.TextUtils;
 import android.util.Log;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
-import java.util.ArrayList;
-import java.util.List;
+import java.util.HashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
@@ -51,11 +50,12 @@ 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;
-    private static final int ADD_CLIENT                 = 3;
-    private static final int REMOVE_CLIENT              = 4;
 
     private boolean mEnableTrafficStatsPoll = false;
     private int mTrafficStatsPollToken = 0;
@@ -64,7 +64,7 @@ public class WifiTrafficPoller {
     /* Tracks last reported data activity */
     private int mDataActivity;
 
-    private final List<Messenger> mClients = new ArrayList<Messenger>();
+    private final HashMap<Integer, ITrafficStateCallback> mRegisteredCallbacks = new HashMap<>();
     // err on the side of updating at boot since screen on broadcast may be missed
     // the first time
     private AtomicBoolean mScreenOn = new AtomicBoolean(true);
@@ -105,14 +105,29 @@ public class WifiTrafficPoller {
                 }, filter);
     }
 
-    /** */
-    public void addClient(Messenger client) {
-        Message.obtain(mTrafficHandler, ADD_CLIENT, client).sendToTarget();
+    /**
+     * 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());
+        }
+        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());
+        }
     }
 
-    /** */
-    public void removeClient(Messenger client) {
-        Message.obtain(mTrafficHandler, REMOVE_CLIENT, client).sendToTarget();
+    /**
+     * Remove an existing callback to the traffic poller.
+     */
+    public void removeCallback(int callbackIdentifier) {
+        mRegisteredCallbacks.remove(callbackIdentifier);
+        if (mVerboseLoggingEnabled) {
+            Log.d(TAG, "Removing callback. Num callbacks: " + mRegisteredCallbacks.size());
+        }
     }
 
     void enableVerboseLogging(int verbose) {
@@ -151,7 +166,7 @@ public class WifiTrafficPoller {
                         Log.d(TAG, "TRAFFIC_STATS_POLL "
                                 + mEnableTrafficStatsPoll + " Token "
                                 + Integer.toString(mTrafficStatsPollToken)
-                                + " num clients " + mClients.size());
+                                + " num clients " + mRegisteredCallbacks.size());
                     }
                     if (msg.arg1 == mTrafficStatsPollToken) {
                         ifaceName = mWifiNative.getClientInterfaceName();
@@ -162,18 +177,7 @@ public class WifiTrafficPoller {
                         }
                     }
                     break;
-                case ADD_CLIENT:
-                    mClients.add((Messenger) msg.obj);
-                    if (mVerboseLoggingEnabled) {
-                        Log.d(TAG, "ADD_CLIENT: "
-                                + Integer.toString(mClients.size()));
-                    }
-                    break;
-                case REMOVE_CLIENT:
-                    mClients.remove(msg.obj);
-                    break;
             }
-
         }
     }
 
@@ -193,8 +197,9 @@ public class WifiTrafficPoller {
     private void notifyOnDataActivity(@NonNull String ifaceName) {
         long sent, received;
         long preTxPkts = mTxPkts, preRxPkts = mRxPkts;
-        int dataActivity = WifiManager.DATA_ACTIVITY_NONE;
+        int dataActivity = WifiManager.TrafficStateCallback.DATA_ACTIVITY_NONE;
 
+        // TODO (b/111691443): Use WifiInfo instead of making the native calls here.
         mTxPkts = mWifiNative.getTxPackets(ifaceName);
         mRxPkts = mWifiNative.getRxPackets(ifaceName);
 
@@ -209,10 +214,10 @@ public class WifiTrafficPoller {
             sent = mTxPkts - preTxPkts;
             received = mRxPkts - preRxPkts;
             if (sent > 0) {
-                dataActivity |= WifiManager.DATA_ACTIVITY_OUT;
+                dataActivity |= WifiManager.TrafficStateCallback.DATA_ACTIVITY_OUT;
             }
             if (received > 0) {
-                dataActivity |= WifiManager.DATA_ACTIVITY_IN;
+                dataActivity |= WifiManager.TrafficStateCallback.DATA_ACTIVITY_IN;
             }
 
             if (dataActivity != mDataActivity && mScreenOn.get()) {
@@ -221,12 +226,9 @@ public class WifiTrafficPoller {
                     Log.e(TAG, "notifying of data activity "
                             + Integer.toString(mDataActivity));
                 }
-                for (Messenger client : mClients) {
-                    Message msg = Message.obtain();
-                    msg.what = WifiManager.DATA_ACTIVITY_NOTIFICATION;
-                    msg.arg1 = mDataActivity;
+                for (ITrafficStateCallback callback : mRegisteredCallbacks.values()) {
                     try {
-                        client.send(msg);
+                        callback.onStateChanged(mDataActivity);
                     } catch (RemoteException e) {
                         // Failed to reach, skip
                         // Client removal is handled in WifiService
@@ -242,6 +244,7 @@ public class WifiTrafficPoller {
         pw.println("mTxPkts " + mTxPkts);
         pw.println("mRxPkts " + mRxPkts);
         pw.println("mDataActivity " + mDataActivity);
+        pw.println("mRegisteredCallbacks " + mRegisteredCallbacks);
     }
 
 }
diff --git a/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java b/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java
index addaf354b..f4a97c2b0 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java
@@ -66,6 +66,7 @@ import android.content.pm.PackageManager;
 import android.content.res.Resources;
 import android.net.Uri;
 import android.net.wifi.ISoftApCallback;
+import android.net.wifi.ITrafficStateCallback;
 import android.net.wifi.ScanResult;
 import android.net.wifi.WifiConfiguration;
 import android.net.wifi.WifiConfiguration.KeyMgmt;
@@ -137,6 +138,7 @@ public class WifiServiceImplTest {
     private static final int TEST_UID = 1200000;
     private static final int OTHER_TEST_UID = 1300000;
     private static final int TEST_USER_HANDLE = 13;
+    private static final int TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER = 17;
     private static final String WIFI_IFACE_NAME = "wlan0";
     private static final String TEST_COUNTRY_CODE = "US";
 
@@ -197,6 +199,7 @@ public class WifiServiceImplTest {
     @Mock PowerProfile mPowerProfile;
     @Mock WifiTrafficPoller mWifiTrafficPolller;
     @Mock ScanRequestProxy mScanRequestProxy;
+    @Mock ITrafficStateCallback mTrafficStateCallback;
 
     @Spy FakeWifiLog mLog;
 
@@ -2865,4 +2868,101 @@ public class WifiServiceImplTest {
             fail("Expected Security exception");
         } catch (SecurityException e) { }
     }
+
+    /**
+     * Verify that a call to registerTrafficStateCallback throws a SecurityException if the caller
+     * does not have NETWORK_SETTINGS permission.
+     */
+    @Test
+    public void registerTrafficStateCallbackThrowsSecurityExceptionOnMissingPermissions() {
+        doThrow(new SecurityException()).when(mContext)
+                .enforceCallingOrSelfPermission(eq(android.Manifest.permission.NETWORK_SETTINGS),
+                        eq("WifiService"));
+        try {
+            mWifiServiceImpl.registerTrafficStateCallback(mAppBinder, mTrafficStateCallback,
+                    TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+            fail("expected SecurityException");
+        } catch (SecurityException expected) {
+        }
+    }
+
+    /**
+     * Verify that a call to registerTrafficStateCallback throws an IllegalArgumentException if the
+     * parameters are not provided.
+     */
+    @Test
+    public void registerTrafficStateCallbackThrowsIllegalArgumentExceptionOnInvalidArguments() {
+        try {
+            mWifiServiceImpl.registerTrafficStateCallback(
+                    mAppBinder, null, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+            fail("expected IllegalArgumentException");
+        } catch (IllegalArgumentException expected) {
+        }
+    }
+
+    /**
+     * Verify that a call to unregisterTrafficStateCallback throws a SecurityException if the caller
+     * does not have NETWORK_SETTINGS permission.
+     */
+    @Test
+    public void unregisterTrafficStateCallbackThrowsSecurityExceptionOnMissingPermissions() {
+        doThrow(new SecurityException()).when(mContext)
+                .enforceCallingOrSelfPermission(eq(android.Manifest.permission.NETWORK_SETTINGS),
+                        eq("WifiService"));
+        try {
+            mWifiServiceImpl.unregisterTrafficStateCallback(TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+            fail("expected SecurityException");
+        } catch (SecurityException expected) {
+        }
+    }
+
+    /**
+     * Verify that registerTrafficStateCallback adds callback to {@link WifiTrafficPoller}.
+     */
+    @Test
+    public void registerTrafficStateCallbackAndVerify() throws Exception {
+        mWifiServiceImpl.registerTrafficStateCallback(
+                mAppBinder, mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+        mLooper.dispatchAll();
+        verify(mWifiTrafficPoller).addCallback(
+                mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+    }
+
+    /**
+     * Verify that unregisterTrafficStateCallback removes callback from {@link WifiTrafficPoller}.
+     */
+    @Test
+    public void unregisterTrafficStateCallbackAndVerify() throws Exception {
+        mWifiServiceImpl.unregisterTrafficStateCallback(0);
+        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 40281f8ba..b00c4f74c 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiTrafficPollerTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiTrafficPollerTest.java
@@ -15,32 +15,31 @@
  */
 package com.android.server.wifi;
 
-import static org.junit.Assert.assertEquals;
 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.never;
-import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import org.junit.Before;
-import org.junit.Test;
-import org.mockito.ArgumentCaptor;
-import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
-
 import android.content.BroadcastReceiver;
 import android.content.Context;
 import android.content.Intent;
 import android.net.NetworkInfo;
+import android.net.wifi.ITrafficStateCallback;
 import android.net.wifi.WifiManager;
-import android.os.Handler;
 import android.os.Message;
-import android.os.Messenger;
+import android.os.RemoteException;
 import android.os.test.TestLooper;
 import android.support.test.filters.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.WifiTrafficPoller}.
  */
@@ -49,15 +48,14 @@ public class WifiTrafficPollerTest {
     public static final String TAG = "WifiTrafficPollerTest";
 
     private TestLooper mLooper;
-    private Handler mHandler;
     private WifiTrafficPoller mWifiTrafficPoller;
     private BroadcastReceiver mReceiver;
     private Intent mIntent;
-    private Messenger mMessenger;
     private final static String IFNAME = "wlan0";
     private final static long DEFAULT_PACKET_COUNT = 10;
     private final static long TX_PACKET_COUNT = 40;
     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 =
@@ -66,6 +64,7 @@ public class WifiTrafficPollerTest {
     @Mock Context mContext;
     @Mock WifiNative mWifiNative;
     @Mock NetworkInfo mNetworkInfo;
+    @Mock ITrafficStateCallback mTrafficStateCallback;
 
     /**
      * Called before each test
@@ -74,8 +73,6 @@ public class WifiTrafficPollerTest {
     public void setUp() throws Exception {
         // Ensure looper exists
         mLooper = new TestLooper();
-        mHandler = spy(new Handler(mLooper.getLooper()));
-        mMessenger = new Messenger(mHandler);
         MockitoAnnotations.initMocks(this);
 
         when(mWifiNative.getTxPackets(any(String.class))).thenReturn(DEFAULT_PACKET_COUNT,
@@ -98,12 +95,6 @@ public class WifiTrafficPollerTest {
                 NetworkInfo.DetailedState.CONNECTED);
     }
 
-    private void registerClient() {
-        // Register Client to verify that Tx/RX packet message is properly received.
-        mWifiTrafficPoller.addClient(mMessenger);
-        mLooper.dispatchAll();
-    }
-
     private void triggerForUpdatedInformationOfData(String actionScreen,
             NetworkInfo.DetailedState networkState) {
         when(mNetworkInfo.getDetailedState()).thenReturn(NetworkInfo.DetailedState.DISCONNECTED);
@@ -122,13 +113,15 @@ public class WifiTrafficPollerTest {
      * Verify that StartTrafficStatsPolling should not happen in case a network is not connected
      */
     @Test
-    public void testNotStartTrafficStatsPollingWithDisconnected() {
-        registerClient();
+    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);
         triggerForUpdatedInformationOfData(Intent.ACTION_SCREEN_ON,
                 NetworkInfo.DetailedState.DISCONNECTED);
 
         // Client should not get any message when the network is disconnected
-        verify(mHandler, never()).handleMessage(any(Message.class));
+        verify(mTrafficStateCallback, never()).onStateChanged(anyInt());
     }
 
     /**
@@ -136,22 +129,26 @@ public class WifiTrafficPollerTest {
      * available.
      */
     @Test
-    public void testStartTrafficStatsPollingWithScreenOn() {
-        registerClient();
+    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);
         triggerForUpdatedInformationOfData(Intent.ACTION_SCREEN_ON,
                 NetworkInfo.DetailedState.CONNECTED);
 
         // Client should get the DATA_ACTIVITY_NOTIFICATION
-        verify(mHandler).handleMessage(mMessageCaptor.capture());
-        assertEquals(WifiManager.DATA_ACTIVITY_NOTIFICATION, mMessageCaptor.getValue().what);
+        verify(mTrafficStateCallback).onStateChanged(
+                WifiManager.TrafficStateCallback.DATA_ACTIVITY_INOUT);
     }
 
     /**
      * Verify that StartTrafficStatsPolling should not happen in case screen is off.
      */
     @Test
-    public void testNotStartTrafficStatsPollingWithScreenOff() {
-        registerClient();
+    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);
         triggerForUpdatedInformationOfData(Intent.ACTION_SCREEN_OFF,
                 NetworkInfo.DetailedState.CONNECTED);
 
@@ -159,22 +156,43 @@ public class WifiTrafficPollerTest {
         mLooper.dispatchAll();
 
         // Client should not get any message when the screen is off
-        verify(mHandler, never()).handleMessage(any(Message.class));
+        verify(mTrafficStateCallback, never()).onStateChanged(anyInt());
     }
 
     /**
-     * Verify that remove client message should be handled
+     * Verify that remove client should be handled
      */
     @Test
-    public void testRemoveClient() {
-        registerClient();
-        mWifiTrafficPoller.removeClient(mMessenger);
+    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);
+        mWifiTrafficPoller.removeCallback(TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
         mLooper.dispatchAll();
 
         triggerForUpdatedInformationOfData(Intent.ACTION_SCREEN_ON,
                 NetworkInfo.DetailedState.CONNECTED);
 
         // Client should not get any message after the client is removed.
-        verify(mHandler, never()).handleMessage(any(Message.class));
+        verify(mTrafficStateCallback, never()).onStateChanged(anyInt());
+    }
+
+    /**
+     * Verify that remove client ignores when callback identifier is wrong.
+     */
+    @Test
+    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);
+        mWifiTrafficPoller.removeCallback(TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER + 5);
+        mLooper.dispatchAll();
+
+        triggerForUpdatedInformationOfData(Intent.ACTION_SCREEN_ON,
+                NetworkInfo.DetailedState.CONNECTED);
+
+        // Client should get the DATA_ACTIVITY_NOTIFICATION
+        verify(mTrafficStateCallback).onStateChanged(
+                WifiManager.TrafficStateCallback.DATA_ACTIVITY_INOUT);
     }
 }