From d629e84dbdd13640f0b183053bca7a4ca9672826 Mon Sep 17 00:00:00 2001
From: Junyu Lai <junyulai@google.com>
Date: Wed, 25 May 2022 10:37:09 +0800
Subject: [PATCH] Fix crash when incrementing data usage request

This fixed a request leak problem which was introduced in previous
Android release but caught recently, where the budget of
MultipathPolicyTracker was cleaned before unregistering requests,
the logic in the tracker would skip the unregister if there is
no budget.

Test: MultipathPolicyTrackerTest#testOnThresholdReached
Bug: 233733267
Fix: 145963406
Change-Id: Icee43a4ee913a6459a47632f608fa6cb8b2ed4c2
Merged-In: Icee43a4ee913a6459a47632f608fa6cb8b2ed4c2
  (pure cherry-picked from ag/18583185)
---
 .../connectivity/MultipathPolicyTracker.java  | 33 ++++++++++++-------
 1 file changed, 21 insertions(+), 12 deletions(-)

diff --git a/services/core/java/com/android/server/connectivity/MultipathPolicyTracker.java b/services/core/java/com/android/server/connectivity/MultipathPolicyTracker.java
index b9efdf5516464..574a3e7eab7eb 100644
--- a/services/core/java/com/android/server/connectivity/MultipathPolicyTracker.java
+++ b/services/core/java/com/android/server/connectivity/MultipathPolicyTracker.java
@@ -192,10 +192,13 @@ public class MultipathPolicyTracker {
         private final int mSubId;
 
         private long mQuota;
-        /** Current multipath budget. Nonzero iff we have budget and a UsageCallback is armed. */
-        private long mMultipathBudget;
+        /** Current multipath budget. Nonzero iff we have budget. */
+        // The budget could be accessed by multiple threads, make it volatile to ensure the callers
+        // on a different thread will not see the stale value.
+        private volatile long mMultipathBudget;
         private final NetworkTemplate mNetworkTemplate;
         private final UsageCallback mUsageCallback;
+        private boolean mUsageCallbackRegistered = false;
         private NetworkCapabilities mNetworkCapabilities;
         private final NetworkStatsManager mStatsManager;
 
@@ -234,7 +237,6 @@ public class MultipathPolicyTracker {
                 @Override
                 public void onThresholdReached(int networkType, String subscriberId) {
                     if (DBG) Log.d(TAG, "onThresholdReached for network " + network);
-                    mMultipathBudget = 0;
                     updateMultipathBudget();
                 }
             };
@@ -376,9 +378,9 @@ public class MultipathPolicyTracker {
                 if (DBG) {
                     Log.d(TAG, "Setting callback for " + budget + " bytes on network " + network);
                 }
-                registerUsageCallback(budget);
+                setMultipathBudget(budget);
             } else {
-                maybeUnregisterUsageCallback();
+                clearMultipathBudget();
             }
         }
 
@@ -403,23 +405,30 @@ public class MultipathPolicyTracker {
             return mMultipathBudget > 0;
         }
 
-        private void registerUsageCallback(long budget) {
+        // Sets the budget and registers a usage callback for it.
+        private void setMultipathBudget(long budget) {
             maybeUnregisterUsageCallback();
+            if (DBG) Log.d(TAG, "Registering callback, budget is " + mMultipathBudget);
             mStatsManager.registerUsageCallback(mNetworkTemplate, budget,
                     (command) -> mHandler.post(command), mUsageCallback);
+            mUsageCallbackRegistered = true;
             mMultipathBudget = budget;
         }
 
         private void maybeUnregisterUsageCallback() {
-            if (haveMultipathBudget()) {
-                if (DBG) Log.d(TAG, "Unregistering callback, budget was " + mMultipathBudget);
-                mStatsManager.unregisterUsageCallback(mUsageCallback);
-                mMultipathBudget = 0;
-            }
+            if (!mUsageCallbackRegistered) return;
+            if (DBG) Log.d(TAG, "Unregistering callback, budget was " + mMultipathBudget);
+            mStatsManager.unregisterUsageCallback(mUsageCallback);
+            mUsageCallbackRegistered = false;
+        }
+
+        private void clearMultipathBudget() {
+            maybeUnregisterUsageCallback();
+            mMultipathBudget = 0;
         }
 
         void shutdown() {
-            maybeUnregisterUsageCallback();
+            clearMultipathBudget();
         }
     }