Merge "[Thread] start / initialize Thread service early" into main
diff --git a/DnsResolver/include/DnsHelperPublic.h b/DnsResolver/include/DnsHelperPublic.h
index 7c9fc9e..44b0012 100644
--- a/DnsResolver/include/DnsHelperPublic.h
+++ b/DnsResolver/include/DnsHelperPublic.h
@@ -25,7 +25,8 @@
  * Perform any required initialization - including opening any required BPF maps. This function
  * needs to be called before using other functions of this library.
  *
- * Returns 0 on success, a negative POSIX error code (see errno.h) on other failures.
+ * Returns 0 on success, -EOPNOTSUPP when the function is called on the Android version before
+ * T. Returns a negative POSIX error code (see errno.h) on other failures.
  */
 int ADnsHelper_init();
 
@@ -36,7 +37,9 @@
  * |uid| is a Linux/Android UID to be queried. It is a combination of UserID and AppID.
  * |metered| indicates whether the uid is currently using a billing network.
  *
- * Returns 0(false)/1(true) on success, a negative POSIX error code (see errno.h) on other failures.
+ * Returns 0(false)/1(true) on success, -EUNATCH when the ADnsHelper_init is not called before
+ * calling this function. Returns a negative POSIX error code (see errno.h) on other failures
+ * that return from bpf syscall.
  */
 int ADnsHelper_isUidNetworkingBlocked(uid_t uid, bool metered);
 
diff --git a/service-t/jni/com_android_server_net_NetworkStatsService.cpp b/service-t/jni/com_android_server_net_NetworkStatsService.cpp
index bdbb655..81912ae 100644
--- a/service-t/jni/com_android_server_net_NetworkStatsService.cpp
+++ b/service-t/jni/com_android_server_net_NetworkStatsService.cpp
@@ -34,77 +34,64 @@
 
 using android::bpf::bpfGetUidStats;
 using android::bpf::bpfGetIfaceStats;
-using android::bpf::bpfGetIfIndexStats;
 using android::bpf::NetworkTraceHandler;
 
 namespace android {
 
-// NOTE: keep these in sync with TrafficStats.java
-static const uint64_t UNKNOWN = -1;
-
-enum StatsType {
-    RX_BYTES = 0,
-    RX_PACKETS = 1,
-    TX_BYTES = 2,
-    TX_PACKETS = 3,
-};
-
-static uint64_t getStatsType(StatsValue* stats, StatsType type) {
-    switch (type) {
-        case RX_BYTES:
-            return stats->rxBytes;
-        case RX_PACKETS:
-            return stats->rxPackets;
-        case TX_BYTES:
-            return stats->txBytes;
-        case TX_PACKETS:
-            return stats->txPackets;
-        default:
-            return UNKNOWN;
+static jobject statsValueToEntry(JNIEnv* env, StatsValue* stats) {
+    // Find the Java class that represents the structure
+    jclass gEntryClass = env->FindClass("android/net/NetworkStats$Entry");
+    if (gEntryClass == nullptr) {
+        return nullptr;
     }
+
+    // Create a new instance of the Java class
+    jobject result = env->AllocObject(gEntryClass);
+    if (result == nullptr) {
+        return nullptr;
+    }
+
+    // Set the values of the structure fields in the Java object
+    env->SetLongField(result, env->GetFieldID(gEntryClass, "rxBytes", "J"), stats->rxBytes);
+    env->SetLongField(result, env->GetFieldID(gEntryClass, "txBytes", "J"), stats->txBytes);
+    env->SetLongField(result, env->GetFieldID(gEntryClass, "rxPackets", "J"), stats->rxPackets);
+    env->SetLongField(result, env->GetFieldID(gEntryClass, "txPackets", "J"), stats->txPackets);
+
+    return result;
 }
 
-static jlong nativeGetTotalStat(JNIEnv* env, jclass clazz, jint type) {
+static jobject nativeGetTotalStat(JNIEnv* env, jclass clazz) {
     StatsValue stats = {};
 
     if (bpfGetIfaceStats(NULL, &stats) == 0) {
-        return getStatsType(&stats, (StatsType) type);
+        return statsValueToEntry(env, &stats);
     } else {
-        return UNKNOWN;
+        return nullptr;
     }
 }
 
-static jlong nativeGetIfaceStat(JNIEnv* env, jclass clazz, jstring iface, jint type) {
+static jobject nativeGetIfaceStat(JNIEnv* env, jclass clazz, jstring iface) {
     ScopedUtfChars iface8(env, iface);
     if (iface8.c_str() == NULL) {
-        return UNKNOWN;
+        return nullptr;
     }
 
     StatsValue stats = {};
 
     if (bpfGetIfaceStats(iface8.c_str(), &stats) == 0) {
-        return getStatsType(&stats, (StatsType) type);
+        return statsValueToEntry(env, &stats);
     } else {
-        return UNKNOWN;
+        return nullptr;
     }
 }
 
-static jlong nativeGetIfIndexStat(JNIEnv* env, jclass clazz, jint ifindex, jint type) {
-    StatsValue stats = {};
-    if (bpfGetIfIndexStats(ifindex, &stats) == 0) {
-        return getStatsType(&stats, (StatsType) type);
-    } else {
-        return UNKNOWN;
-    }
-}
-
-static jlong nativeGetUidStat(JNIEnv* env, jclass clazz, jint uid, jint type) {
+static jobject nativeGetUidStat(JNIEnv* env, jclass clazz, jint uid) {
     StatsValue stats = {};
 
     if (bpfGetUidStats(uid, &stats) == 0) {
-        return getStatsType(&stats, (StatsType) type);
+        return statsValueToEntry(env, &stats);
     } else {
-        return UNKNOWN;
+        return nullptr;
     }
 }
 
@@ -113,11 +100,26 @@
 }
 
 static const JNINativeMethod gMethods[] = {
-        {"nativeGetTotalStat", "(I)J", (void*)nativeGetTotalStat},
-        {"nativeGetIfaceStat", "(Ljava/lang/String;I)J", (void*)nativeGetIfaceStat},
-        {"nativeGetIfIndexStat", "(II)J", (void*)nativeGetIfIndexStat},
-        {"nativeGetUidStat", "(II)J", (void*)nativeGetUidStat},
-        {"nativeInitNetworkTracing", "()V", (void*)nativeInitNetworkTracing},
+        {
+            "nativeGetTotalStat",
+            "()Landroid/net/NetworkStats$Entry;",
+            (void*)nativeGetTotalStat
+        },
+        {
+            "nativeGetIfaceStat",
+            "(Ljava/lang/String;)Landroid/net/NetworkStats$Entry;",
+            (void*)nativeGetIfaceStat
+        },
+        {
+            "nativeGetUidStat",
+            "(I)Landroid/net/NetworkStats$Entry;",
+            (void*)nativeGetUidStat
+        },
+        {
+            "nativeInitNetworkTracing",
+            "()V",
+            (void*)nativeInitNetworkTracing
+        },
 };
 
 int register_android_server_net_NetworkStatsService(JNIEnv* env) {
diff --git a/service-t/src/com/android/server/net/NetworkStatsService.java b/service-t/src/com/android/server/net/NetworkStatsService.java
index 3ac5e29..eb75461 100644
--- a/service-t/src/com/android/server/net/NetworkStatsService.java
+++ b/service-t/src/com/android/server/net/NetworkStatsService.java
@@ -1981,36 +1981,56 @@
         if (callingUid != android.os.Process.SYSTEM_UID && callingUid != uid) {
             return UNSUPPORTED;
         }
-        return nativeGetUidStat(uid, type);
+        return getEntryValueForType(nativeGetUidStat(uid), type);
     }
 
     @Override
     public long getIfaceStats(@NonNull String iface, int type) {
         Objects.requireNonNull(iface);
-        long nativeIfaceStats = nativeGetIfaceStat(iface, type);
-        if (nativeIfaceStats == -1) {
-            return nativeIfaceStats;
+        final NetworkStats.Entry entry = nativeGetIfaceStat(iface);
+        final long value = getEntryValueForType(entry, type);
+        if (value == UNSUPPORTED) {
+            return UNSUPPORTED;
         } else {
             // When tethering offload is in use, nativeIfaceStats does not contain usage from
             // offload, add it back here. Note that the included statistics might be stale
             // since polling newest stats from hardware might impact system health and not
             // suitable for TrafficStats API use cases.
-            return nativeIfaceStats + getProviderIfaceStats(iface, type);
+            entry.add(getProviderIfaceStats(iface));
+            return getEntryValueForType(entry, type);
+        }
+    }
+
+    private long getEntryValueForType(@Nullable NetworkStats.Entry entry, int type) {
+        if (entry == null) return UNSUPPORTED;
+        switch (type) {
+            case TrafficStats.TYPE_RX_BYTES:
+                return entry.rxBytes;
+            case TrafficStats.TYPE_TX_BYTES:
+                return entry.txBytes;
+            case TrafficStats.TYPE_RX_PACKETS:
+                return entry.rxPackets;
+            case TrafficStats.TYPE_TX_PACKETS:
+                return entry.txPackets;
+            default:
+                return UNSUPPORTED;
         }
     }
 
     @Override
     public long getTotalStats(int type) {
-        long nativeTotalStats = nativeGetTotalStat(type);
-        if (nativeTotalStats == -1) {
-            return nativeTotalStats;
+        final NetworkStats.Entry entry = nativeGetTotalStat();
+        final long value = getEntryValueForType(entry, type);
+        if (value == UNSUPPORTED) {
+            return UNSUPPORTED;
         } else {
             // Refer to comment in getIfaceStats
-            return nativeTotalStats + getProviderIfaceStats(IFACE_ALL, type);
+            entry.add(getProviderIfaceStats(IFACE_ALL));
+            return getEntryValueForType(entry, type);
         }
     }
 
-    private long getProviderIfaceStats(@Nullable String iface, int type) {
+    private NetworkStats.Entry getProviderIfaceStats(@Nullable String iface) {
         final NetworkStats providerSnapshot = getNetworkStatsFromProviders(STATS_PER_IFACE);
         final HashSet<String> limitIfaces;
         if (iface == IFACE_ALL) {
@@ -2019,19 +2039,7 @@
             limitIfaces = new HashSet<>();
             limitIfaces.add(iface);
         }
-        final NetworkStats.Entry entry = providerSnapshot.getTotal(null, limitIfaces);
-        switch (type) {
-            case TrafficStats.TYPE_RX_BYTES:
-                return entry.rxBytes;
-            case TrafficStats.TYPE_RX_PACKETS:
-                return entry.rxPackets;
-            case TrafficStats.TYPE_TX_BYTES:
-                return entry.txBytes;
-            case TrafficStats.TYPE_TX_PACKETS:
-                return entry.txPackets;
-            default:
-                return 0;
-        }
+        return providerSnapshot.getTotal(null, limitIfaces);
     }
 
     /**
@@ -3398,10 +3406,13 @@
         }
     }
 
-    private static native long nativeGetTotalStat(int type);
-    private static native long nativeGetIfaceStat(String iface, int type);
-    private static native long nativeGetIfIndexStat(int ifindex, int type);
-    private static native long nativeGetUidStat(int uid, int type);
+    // TODO: Read stats by using BpfNetMapsReader.
+    @Nullable
+    private static native NetworkStats.Entry nativeGetTotalStat();
+    @Nullable
+    private static native NetworkStats.Entry nativeGetIfaceStat(String iface);
+    @Nullable
+    private static native NetworkStats.Entry nativeGetUidStat(int uid);
 
     /** Initializes and registers the Perfetto Network Trace data source */
     public static native void nativeInitNetworkTracing();
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityDiagnosticsManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityDiagnosticsManagerTest.java
index e0fe929..ceb48d4 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityDiagnosticsManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityDiagnosticsManagerTest.java
@@ -298,17 +298,6 @@
                 },
                 android.Manifest.permission.MODIFY_PHONE_STATE);
 
-        // TODO(b/157779832): This should use android.permission.CHANGE_NETWORK_STATE. However, the
-        // shell does not have CHANGE_NETWORK_STATE, so use CONNECTIVITY_INTERNAL until the shell
-        // permissions are updated.
-        runWithShellPermissionIdentity(
-                () -> mConnectivityManager.requestNetwork(
-                        CELLULAR_NETWORK_REQUEST, testNetworkCallback),
-                android.Manifest.permission.CONNECTIVITY_INTERNAL);
-
-        final Network network = testNetworkCallback.waitForAvailable();
-        assertNotNull(network);
-
         assertTrue("Didn't receive broadcast for ACTION_CARRIER_CONFIG_CHANGED for subId=" + subId,
                 carrierConfigReceiver.waitForCarrierConfigChanged());
 
@@ -324,6 +313,17 @@
 
         Thread.sleep(5_000);
 
+        // TODO(b/157779832): This should use android.permission.CHANGE_NETWORK_STATE. However, the
+        // shell does not have CHANGE_NETWORK_STATE, so use CONNECTIVITY_INTERNAL until the shell
+        // permissions are updated.
+        runWithShellPermissionIdentity(
+                () -> mConnectivityManager.requestNetwork(
+                        CELLULAR_NETWORK_REQUEST, testNetworkCallback),
+                android.Manifest.permission.CONNECTIVITY_INTERNAL);
+
+        final Network network = testNetworkCallback.waitForAvailable();
+        assertNotNull(network);
+
         // TODO(b/217559768): Receiving carrier config change and immediately checking carrier
         //  privileges is racy, as the CP status is updated after receiving the same signal. Move
         //  the CP check after sleep to temporarily reduce the flakiness. This will soon be fixed
diff --git a/thread/tests/cts/src/android/net/thread/cts/ThreadNetworkControllerTest.java b/thread/tests/cts/src/android/net/thread/cts/ThreadNetworkControllerTest.java
index 362ff39..e02e74d 100644
--- a/thread/tests/cts/src/android/net/thread/cts/ThreadNetworkControllerTest.java
+++ b/thread/tests/cts/src/android/net/thread/cts/ThreadNetworkControllerTest.java
@@ -521,7 +521,7 @@
     }
 
     @Test
-    public void scheduleMigration_withPrivilegedPermission_success() throws Exception {
+    public void scheduleMigration_withPrivilegedPermission_newDatasetApplied() throws Exception {
         grantPermissions(permission.ACCESS_NETWORK_STATE, PERMISSION_THREAD_NETWORK_PRIVILEGED);
 
         for (ThreadNetworkController controller : getAllControllers()) {
@@ -548,11 +548,32 @@
 
             controller.scheduleMigration(
                     pendingDataset2, mExecutor, newOutcomeReceiver(migrateFuture));
-
             migrateFuture.get();
-            Thread.sleep(35 * 1000);
-            assertThat(getActiveOperationalDataset(controller)).isEqualTo(activeDataset2);
-            assertThat(getPendingOperationalDataset(controller)).isNull();
+
+            SettableFuture<Boolean> dataset2IsApplied = SettableFuture.create();
+            SettableFuture<Boolean> pendingDatasetIsRemoved = SettableFuture.create();
+            OperationalDatasetCallback datasetCallback =
+                    new OperationalDatasetCallback() {
+                        @Override
+                        public void onActiveOperationalDatasetChanged(
+                                ActiveOperationalDataset activeDataset) {
+                            if (activeDataset.equals(activeDataset2)) {
+                                dataset2IsApplied.set(true);
+                            }
+                        }
+
+                        @Override
+                        public void onPendingOperationalDatasetChanged(
+                                PendingOperationalDataset pendingDataset) {
+                            if (pendingDataset == null) {
+                                pendingDatasetIsRemoved.set(true);
+                            }
+                        }
+                    };
+            controller.registerOperationalDatasetCallback(directExecutor(), datasetCallback);
+            assertThat(dataset2IsApplied.get()).isTrue();
+            assertThat(pendingDatasetIsRemoved.get()).isTrue();
+            controller.unregisterOperationalDatasetCallback(datasetCallback);
         }
     }
 
@@ -629,7 +650,8 @@
     }
 
     @Test
-    public void scheduleMigration_secondRequestHasLargerTimestamp_success() throws Exception {
+    public void scheduleMigration_secondRequestHasLargerTimestamp_newDatasetApplied()
+            throws Exception {
         grantPermissions(permission.ACCESS_NETWORK_STATE, PERMISSION_THREAD_NETWORK_PRIVILEGED);
 
         for (ThreadNetworkController controller : getAllControllers()) {
@@ -669,11 +691,32 @@
             migrateFuture1.get();
             controller.scheduleMigration(
                     pendingDataset2, mExecutor, newOutcomeReceiver(migrateFuture2));
-
             migrateFuture2.get();
-            Thread.sleep(35 * 1000);
-            assertThat(getActiveOperationalDataset(controller)).isEqualTo(activeDataset2);
-            assertThat(getPendingOperationalDataset(controller)).isNull();
+
+            SettableFuture<Boolean> dataset2IsApplied = SettableFuture.create();
+            SettableFuture<Boolean> pendingDatasetIsRemoved = SettableFuture.create();
+            OperationalDatasetCallback datasetCallback =
+                    new OperationalDatasetCallback() {
+                        @Override
+                        public void onActiveOperationalDatasetChanged(
+                                ActiveOperationalDataset activeDataset) {
+                            if (activeDataset.equals(activeDataset2)) {
+                                dataset2IsApplied.set(true);
+                            }
+                        }
+
+                        @Override
+                        public void onPendingOperationalDatasetChanged(
+                                PendingOperationalDataset pendingDataset) {
+                            if (pendingDataset == null) {
+                                pendingDatasetIsRemoved.set(true);
+                            }
+                        }
+                    };
+            controller.registerOperationalDatasetCallback(directExecutor(), datasetCallback);
+            assertThat(dataset2IsApplied.get()).isTrue();
+            assertThat(pendingDatasetIsRemoved.get()).isTrue();
+            controller.unregisterOperationalDatasetCallback(datasetCallback);
         }
     }