Merge "Fix flaky test: testNsdManager_registerOffloadEngine" into main
diff --git a/Tethering/src/com/android/networkstack/tethering/OffloadHardwareInterface.java b/Tethering/src/com/android/networkstack/tethering/OffloadHardwareInterface.java
index de15c5b..53c80ae 100644
--- a/Tethering/src/com/android/networkstack/tethering/OffloadHardwareInterface.java
+++ b/Tethering/src/com/android/networkstack/tethering/OffloadHardwareInterface.java
@@ -283,13 +283,16 @@
 
     private int initWithHandles(NativeHandle h1, NativeHandle h2) {
         if (h1 == null || h2 == null) {
+            // Set mIOffload to null has two purposes:
+            // 1. NativeHandles can be closed after initWithHandles() fails
+            // 2. Prevent mIOffload.stopOffload() to be called in stopOffload()
+            mIOffload = null;
             mLog.e("Failed to create socket.");
             return OFFLOAD_HAL_VERSION_NONE;
         }
 
         requestSocketDump(h1);
         if (!mIOffload.initOffload(h1, h2, mOffloadHalCallback)) {
-            mIOffload.stopOffload();
             mLog.e("Failed to initialize offload.");
             return OFFLOAD_HAL_VERSION_NONE;
         }
@@ -329,9 +332,9 @@
         mOffloadHalCallback = offloadCb;
         final int version = initWithHandles(h1, h2);
 
-        // Explicitly close FDs for HIDL. AIDL will pass the original FDs to the service,
-        // they shouldn't be closed here.
-        if (version < OFFLOAD_HAL_VERSION_AIDL) {
+        // Explicitly close FDs for HIDL or when mIOffload is null (cleared in initWithHandles).
+        // AIDL will pass the original FDs to the service, they shouldn't be closed here.
+        if (mIOffload == null || mIOffload.getVersion() < OFFLOAD_HAL_VERSION_AIDL) {
             maybeCloseFdInNativeHandles(h1, h2);
         }
         return version;
diff --git a/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java b/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java
index bcedbef..4594f71 100644
--- a/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java
+++ b/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java
@@ -66,7 +66,8 @@
                     event.getFoundServiceCount(),
                     event.getFoundCallbackCount(),
                     event.getLostCallbackCount(),
-                    event.getRepliedRequestsCount());
+                    event.getRepliedRequestsCount(),
+                    event.getSentQueryCount());
         }
     }
 
diff --git a/service-t/src/com/android/server/NsdService.java b/service-t/src/com/android/server/NsdService.java
index 53782ac..f7edbe4 100644
--- a/service-t/src/com/android/server/NsdService.java
+++ b/service-t/src/com/android/server/NsdService.java
@@ -1920,11 +1920,14 @@
     @Override
     public INsdServiceConnector connect(INsdManagerCallback cb, boolean useJavaBackend) {
         mContext.enforceCallingOrSelfPermission(android.Manifest.permission.INTERNET, "NsdService");
+        final int uid = mDeps.getCallingUid();
+        if (cb == null) {
+            throw new IllegalArgumentException("Unknown client callback from uid=" + uid);
+        }
         if (DBG) Log.d(TAG, "New client connect. useJavaBackend=" + useJavaBackend);
         final INsdServiceConnector connector = new NsdServiceConnector();
         mNsdStateMachine.sendMessage(mNsdStateMachine.obtainMessage(NsdManager.REGISTER_CLIENT,
-                new ConnectorArgs((NsdServiceConnector) connector, cb, useJavaBackend,
-                        mDeps.getCallingUid())));
+                new ConnectorArgs((NsdServiceConnector) connector, cb, useJavaBackend, uid)));
         return connector;
     }
 
diff --git a/service/src/com/android/metrics/stats.proto b/service/src/com/android/metrics/stats.proto
index 006d20a..99afb90 100644
--- a/service/src/com/android/metrics/stats.proto
+++ b/service/src/com/android/metrics/stats.proto
@@ -61,6 +61,9 @@
 
   // Record query service count before unregistered service
   optional int32 replied_requests_count = 11;
+
+  // Record sent query count before stopped discovery
+  optional int32 sent_query_count = 12;
 }
 
 /**
diff --git a/tests/unit/java/com/android/server/NsdServiceTest.java b/tests/unit/java/com/android/server/NsdServiceTest.java
index 55384b3..f778075 100644
--- a/tests/unit/java/com/android/server/NsdServiceTest.java
+++ b/tests/unit/java/com/android/server/NsdServiceTest.java
@@ -1585,6 +1585,20 @@
         lockOrder.verify(mMulticastLock).release();
     }
 
+    @Test
+    public void testNullINsdManagerCallback() {
+        final NsdService service = new NsdService(mContext, mHandler, CLEANUP_DELAY_MS, mDeps) {
+            @Override
+            public INsdServiceConnector connect(INsdManagerCallback baseCb,
+                    boolean runNewMdnsBackend) {
+                // Pass null INsdManagerCallback
+                return super.connect(null /* cb */, runNewMdnsBackend);
+            }
+        };
+
+        assertThrows(IllegalArgumentException.class, () -> new NsdManager(mContext, service));
+    }
+
     private void waitForIdle() {
         HandlerUtils.waitForIdle(mHandler, TIMEOUT_MS);
     }