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

2 files changed, 107 insertions, 52 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);
     }
 
 }