[SP29] Send interface warning bytes to NetworkStatsProvider

This change contains necessary modification in NPMS and NSS
to send warning bytes to NetworkStatsProvider. But since
no any provider has been upgraded to handle such parameter.
Thus, no behavior change is made in this patch.

Test: atest NetworkPolicyManagerServiceTest NetworkStatsServiceTest
Test: atest NetworkPolicyManagerServiceTest#testStatsProviderWarningAndLimitReached
Bug: 149467454
Bug: 170699770
Bug: 170179169
Merged-In: I6c4863030c36328db571294fd12a40e59864def5
Change-Id: I6c4863030c36328db571294fd12a40e59864def5
  (cherry-picked from ag/13982166)
diff --git a/services/core/java/com/android/server/net/NetworkPolicyManagerService.java b/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
index 122f761..e9530de 100644
--- a/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
+++ b/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
@@ -383,8 +383,8 @@
     private static final int MSG_LIMIT_REACHED = 5;
     private static final int MSG_RESTRICT_BACKGROUND_CHANGED = 6;
     private static final int MSG_ADVISE_PERSIST_THRESHOLD = 7;
-    private static final int MSG_UPDATE_INTERFACE_QUOTA = 10;
-    private static final int MSG_REMOVE_INTERFACE_QUOTA = 11;
+    private static final int MSG_UPDATE_INTERFACE_QUOTAS = 10;
+    private static final int MSG_REMOVE_INTERFACE_QUOTAS = 11;
     private static final int MSG_POLICIES_CHANGED = 13;
     private static final int MSG_RESET_FIREWALL_RULES_BY_UID = 15;
     private static final int MSG_SUBSCRIPTION_OVERRIDE = 16;
@@ -1910,33 +1910,44 @@
             final boolean hasWarning = policy.warningBytes != LIMIT_DISABLED;
             final boolean hasLimit = policy.limitBytes != LIMIT_DISABLED;
             long limitBytes = Long.MAX_VALUE;
-            if (hasLimit && policy.hasCycle()) {
+            long warningBytes = Long.MAX_VALUE;
+            if ((hasLimit || hasWarning) && policy.hasCycle()) {
                 final Pair<ZonedDateTime, ZonedDateTime> cycle = NetworkPolicyManager
                         .cycleIterator(policy).next();
                 final long start = cycle.first.toInstant().toEpochMilli();
                 final long end = cycle.second.toInstant().toEpochMilli();
                 final long totalBytes = getTotalBytes(policy.template, start, end);
 
-                if (policy.lastLimitSnooze < start) {
+                // If the limit notification is not snoozed, the limit quota needs to be calculated.
+                if (hasLimit && policy.lastLimitSnooze < start) {
                     // remaining "quota" bytes are based on total usage in
                     // current cycle. kernel doesn't like 0-byte rules, so we
                     // set 1-byte quota and disable the radio later.
                     limitBytes = Math.max(1, policy.limitBytes - totalBytes);
                 }
+
+                // If the warning notification was snoozed by user, or the service already knows
+                // it is over warning bytes, doesn't need to calculate warning bytes.
+                if (hasWarning && policy.lastWarningSnooze < start
+                        && !policy.isOverWarning(totalBytes)) {
+                    warningBytes = Math.max(1, policy.warningBytes - totalBytes);
+                }
             }
 
-            if (hasLimit || policy.metered) {
+            if (hasWarning || hasLimit || policy.metered) {
                 if (matchingIfaces.size() > 1) {
                     // TODO: switch to shared quota once NMS supports
                     Slog.w(TAG, "shared quota unsupported; generating rule for each iface");
                 }
 
-                // Set the interface limit. For interfaces which has no cycle, or metered with
-                // no policy limit, or snoozed limit notification; we still need to put iptables
-                // rule hooks to restrict apps for data saver, so push really high quota.
+                // Set the interface warning and limit. For interfaces which has no cycle,
+                // or metered with no policy quotas, or snoozed notification; we still need to put
+                // iptables rule hooks to restrict apps for data saver, so push really high quota.
+                // TODO: Push NetworkStatsProvider.QUOTA_UNLIMITED instead of Long.MAX_VALUE to
+                //  providers.
                 for (int j = matchingIfaces.size() - 1; j >= 0; j--) {
                     final String iface = matchingIfaces.valueAt(j);
-                    setInterfaceQuotaAsync(iface, limitBytes);
+                    setInterfaceQuotasAsync(iface, warningBytes, limitBytes);
                     newMeteredIfaces.add(iface);
                 }
             }
@@ -1960,7 +1971,7 @@
                 for (int j = matchingIfaces.size() - 1; j >= 0; j--) {
                     final String iface = matchingIfaces.valueAt(j);
                     if (!newMeteredIfaces.contains(iface)) {
-                        setInterfaceQuotaAsync(iface, Long.MAX_VALUE);
+                        setInterfaceQuotasAsync(iface, Long.MAX_VALUE, Long.MAX_VALUE);
                         newMeteredIfaces.add(iface);
                     }
                 }
@@ -1972,7 +1983,7 @@
             for (int i = mMeteredIfaces.size() - 1; i >= 0; i--) {
                 final String iface = mMeteredIfaces.valueAt(i);
                 if (!newMeteredIfaces.contains(iface)) {
-                    removeInterfaceQuotaAsync(iface);
+                    removeInterfaceQuotasAsync(iface);
                 }
             }
             mMeteredIfaces = newMeteredIfaces;
@@ -4693,19 +4704,20 @@
                     mNetworkStats.advisePersistThreshold(persistThreshold);
                     return true;
                 }
-                case MSG_UPDATE_INTERFACE_QUOTA: {
-                    final String iface = (String) msg.obj;
-                    // int params need to be stitched back into a long
-                    final long quota = ((long) msg.arg1 << 32) | (msg.arg2 & 0xFFFFFFFFL);
-                    removeInterfaceQuota(iface);
-                    setInterfaceQuota(iface, quota);
-                    mNetworkStats.setStatsProviderLimitAsync(iface, quota);
+                case MSG_UPDATE_INTERFACE_QUOTAS: {
+                    final IfaceQuotas val = (IfaceQuotas) msg.obj;
+                    // TODO: Consider set a new limit before removing the original one.
+                    removeInterfaceLimit(val.iface);
+                    setInterfaceLimit(val.iface, val.limit);
+                    mNetworkStats.setStatsProviderWarningAndLimitAsync(val.iface, val.warning,
+                            val.limit);
                     return true;
                 }
-                case MSG_REMOVE_INTERFACE_QUOTA: {
+                case MSG_REMOVE_INTERFACE_QUOTAS: {
                     final String iface = (String) msg.obj;
-                    removeInterfaceQuota(iface);
-                    mNetworkStats.setStatsProviderLimitAsync(iface, QUOTA_UNLIMITED);
+                    removeInterfaceLimit(iface);
+                    mNetworkStats.setStatsProviderWarningAndLimitAsync(iface, QUOTA_UNLIMITED,
+                            QUOTA_UNLIMITED);
                     return true;
                 }
                 case MSG_RESET_FIREWALL_RULES_BY_UID: {
@@ -4833,15 +4845,32 @@
         }
     }
 
-    private void setInterfaceQuotaAsync(String iface, long quotaBytes) {
-        // long quotaBytes split up into two ints to fit in message
-        mHandler.obtainMessage(MSG_UPDATE_INTERFACE_QUOTA, (int) (quotaBytes >> 32),
-                (int) (quotaBytes & 0xFFFFFFFF), iface).sendToTarget();
+    private static final class IfaceQuotas {
+        @NonNull public final String iface;
+        // Warning and limit bytes of interface qutoas, could be QUOTA_UNLIMITED or Long.MAX_VALUE
+        // if not set. 0 is not acceptable since kernel doesn't like 0-byte rules.
+        public final long warning;
+        public final long limit;
+
+        private IfaceQuotas(@NonNull String iface, long warning, long limit) {
+            this.iface = iface;
+            this.warning = warning;
+            this.limit = limit;
+        }
     }
 
-    private void setInterfaceQuota(String iface, long quotaBytes) {
+    private void setInterfaceQuotasAsync(@NonNull String iface,
+            long warningBytes, long limitBytes) {
+        mHandler.obtainMessage(MSG_UPDATE_INTERFACE_QUOTAS,
+                new IfaceQuotas(iface, warningBytes, limitBytes)).sendToTarget();
+    }
+
+    private void setInterfaceLimit(String iface, long limitBytes) {
         try {
-            mNetworkManager.setInterfaceQuota(iface, quotaBytes);
+            // For legacy design the data warning is covered by global alert, where the
+            // kernel will notify upper layer for a small amount of change of traffic
+            // statistics. Thus, passing warning is not needed.
+            mNetworkManager.setInterfaceQuota(iface, limitBytes);
         } catch (IllegalStateException e) {
             Log.wtf(TAG, "problem setting interface quota", e);
         } catch (RemoteException e) {
@@ -4849,11 +4878,11 @@
         }
     }
 
-    private void removeInterfaceQuotaAsync(String iface) {
-        mHandler.obtainMessage(MSG_REMOVE_INTERFACE_QUOTA, iface).sendToTarget();
+    private void removeInterfaceQuotasAsync(String iface) {
+        mHandler.obtainMessage(MSG_REMOVE_INTERFACE_QUOTAS, iface).sendToTarget();
     }
 
-    private void removeInterfaceQuota(String iface) {
+    private void removeInterfaceLimit(String iface) {
         try {
             mNetworkManager.removeInterfaceQuota(iface);
         } catch (IllegalStateException e) {
diff --git a/services/core/java/com/android/server/net/NetworkStatsManagerInternal.java b/services/core/java/com/android/server/net/NetworkStatsManagerInternal.java
index 0cb0bc2c..0e9a9da 100644
--- a/services/core/java/com/android/server/net/NetworkStatsManagerInternal.java
+++ b/services/core/java/com/android/server/net/NetworkStatsManagerInternal.java
@@ -37,8 +37,9 @@
     public abstract void forceUpdate();
 
     /**
-     * Set the quota limit to all registered custom network stats providers.
+     * Set the warning and limit to all registered custom network stats providers.
      * Note that invocation of any interface will be sent to all providers.
      */
-    public abstract void setStatsProviderLimitAsync(@NonNull String iface, long quota);
+    public abstract void setStatsProviderWarningAndLimitAsync(@NonNull String iface, long warning,
+            long limit);
 }
diff --git a/services/core/java/com/android/server/net/NetworkStatsService.java b/services/core/java/com/android/server/net/NetworkStatsService.java
index 191ff096..0bc2054 100644
--- a/services/core/java/com/android/server/net/NetworkStatsService.java
+++ b/services/core/java/com/android/server/net/NetworkStatsService.java
@@ -1658,11 +1658,14 @@
         }
 
         @Override
-        public void setStatsProviderLimitAsync(@NonNull String iface, long quota) {
-            if (LOGV) Slog.v(TAG, "setStatsProviderLimitAsync(" + iface + "," + quota + ")");
-            // TODO: Set warning accordingly.
+        public void setStatsProviderWarningAndLimitAsync(
+                @NonNull String iface, long warning, long limit) {
+            if (LOGV) {
+                Slog.v(TAG, "setStatsProviderWarningAndLimitAsync("
+                        + iface + "," + warning + "," + limit + ")");
+            }
             invokeForAllStatsProviderCallbacks((cb) -> cb.mProvider.onSetWarningAndLimit(iface,
-                    NetworkStatsProvider.QUOTA_UNLIMITED, quota));
+                    warning, limit));
         }
     }
 
diff --git a/services/tests/servicestests/src/com/android/server/net/NetworkPolicyManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/net/NetworkPolicyManagerServiceTest.java
index e7f8850..3ccc886 100644
--- a/services/tests/servicestests/src/com/android/server/net/NetworkPolicyManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/net/NetworkPolicyManagerServiceTest.java
@@ -1760,57 +1760,75 @@
                 true);
     }
 
-    /**
-     * Test that when StatsProvider triggers limit reached, new limit will be calculated and
-     * re-armed.
-     */
-    @Test
-    public void testStatsProviderLimitReached() throws Exception {
-        final int CYCLE_DAY = 15;
-
-        final NetworkStats stats = new NetworkStats(0L, 1);
+    private void increaseMockedTotalBytes(NetworkStats stats, long rxBytes, long txBytes) {
         stats.insertEntry(TEST_IFACE, UID_A, SET_ALL, TAG_NONE,
-                2999, 1, 2000, 1, 0);
+                rxBytes, 1, txBytes, 1, 0);
         when(mStatsService.getNetworkTotalBytes(any(), anyLong(), anyLong()))
                 .thenReturn(stats.getTotalBytes());
         when(mStatsService.getNetworkUidBytes(any(), anyLong(), anyLong()))
                 .thenReturn(stats);
+    }
+
+    private void triggerOnStatsProviderWarningOrLimitReached() throws InterruptedException {
+        final NetworkPolicyManagerInternal npmi = LocalServices
+                .getService(NetworkPolicyManagerInternal.class);
+        npmi.onStatsProviderWarningOrLimitReached("TEST");
+        // Wait for processing of MSG_STATS_PROVIDER_WARNING_OR_LIMIT_REACHED.
+        postMsgAndWaitForCompletion();
+        verify(mStatsService).forceUpdate();
+        // Wait for processing of MSG_*_INTERFACE_QUOTAS.
+        postMsgAndWaitForCompletion();
+    }
+
+    /**
+     * Test that when StatsProvider triggers warning and limit reached, new quotas will be
+     * calculated and re-armed.
+     */
+    @Test
+    public void testStatsProviderWarningAndLimitReached() throws Exception {
+        final int CYCLE_DAY = 15;
+
+        final NetworkStats stats = new NetworkStats(0L, 1);
+        increaseMockedTotalBytes(stats, 2999, 2000);
 
         // Get active mobile network in place
         expectMobileDefaults();
         mService.updateNetworks();
-        verify(mStatsService).setStatsProviderLimitAsync(TEST_IFACE, Long.MAX_VALUE);
+        verify(mStatsService).setStatsProviderWarningAndLimitAsync(TEST_IFACE, Long.MAX_VALUE,
+                Long.MAX_VALUE);
 
-        // Set limit to 10KB.
+        // Set warning to 7KB and limit to 10KB.
         setNetworkPolicies(new NetworkPolicy(
-                sTemplateMobileAll, CYCLE_DAY, TIMEZONE_UTC, WARNING_DISABLED, 10000L,
-                true));
+                sTemplateMobileAll, CYCLE_DAY, TIMEZONE_UTC, 7000L, 10000L, true));
         postMsgAndWaitForCompletion();
 
-        // Verifies that remaining quota is set to providers.
-        verify(mStatsService).setStatsProviderLimitAsync(TEST_IFACE, 10000L - 4999L);
-
+        // Verifies that remaining quotas are set to providers.
+        verify(mStatsService).setStatsProviderWarningAndLimitAsync(TEST_IFACE, 2001L, 5001L);
         reset(mStatsService);
 
-        // Increase the usage.
-        stats.insertEntry(TEST_IFACE, UID_A, SET_ALL, TAG_NONE,
-                1000, 1, 999, 1, 0);
-        when(mStatsService.getNetworkTotalBytes(any(), anyLong(), anyLong()))
-                .thenReturn(stats.getTotalBytes());
-        when(mStatsService.getNetworkUidBytes(any(), anyLong(), anyLong()))
-                .thenReturn(stats);
+        // Increase the usage and simulates that limit reached fires earlier by provider,
+        // but actually the quota is not yet reached. Verifies that the limit reached leads to
+        // a force update and new quotas should be set.
+        increaseMockedTotalBytes(stats, 1000, 999);
+        triggerOnStatsProviderWarningOrLimitReached();
+        verify(mStatsService).setStatsProviderWarningAndLimitAsync(TEST_IFACE, 2L, 3002L);
+        reset(mStatsService);
 
-        // Simulates that limit reached fires earlier by provider, but actually the quota is not
-        // yet reached.
-        final NetworkPolicyManagerInternal npmi = LocalServices
-                .getService(NetworkPolicyManagerInternal.class);
-        npmi.onStatsProviderWarningOrLimitReached("TEST");
+        // Increase the usage and simulate warning reached, the new warning should be unlimited
+        // since service will disable warning quota to stop lower layer from keep triggering
+        // warning reached event.
+        increaseMockedTotalBytes(stats, 1000L, 1000);
+        triggerOnStatsProviderWarningOrLimitReached();
+        verify(mStatsService).setStatsProviderWarningAndLimitAsync(
+                TEST_IFACE, Long.MAX_VALUE, 1002L);
+        reset(mStatsService);
 
-        // Verifies that the limit reached leads to a force update and new limit should be set.
-        postMsgAndWaitForCompletion();
-        verify(mStatsService).forceUpdate();
-        postMsgAndWaitForCompletion();
-        verify(mStatsService).setStatsProviderLimitAsync(TEST_IFACE, 10000L - 4999L - 1999L);
+        // Increase the usage that over the warning and limit, the new limit should set to 1 to
+        // block the network traffic.
+        increaseMockedTotalBytes(stats, 1000L, 1000);
+        triggerOnStatsProviderWarningOrLimitReached();
+        verify(mStatsService).setStatsProviderWarningAndLimitAsync(TEST_IFACE, Long.MAX_VALUE, 1L);
+        reset(mStatsService);
     }
 
     /**