Support query cancellation for async DNS API

Bug: 124882626
Test: built, flashed, booted
      atest DnsResolverTest DnsPacketTest

Change-Id: Iaa72f5c17f58cf0a58663b892bb18cfdf23cd545

core/java/android/net/DnsResolver.java

@@ -16,6 +16,7 @@
 
 package android.net;
 
+import static android.net.NetworkUtils.resNetworkCancel;
 import static android.net.NetworkUtils.resNetworkQuery;
 import static android.net.NetworkUtils.resNetworkResult;
 import static android.net.NetworkUtils.resNetworkSend;
@@ -26,6 +27,7 @@ import android.annotation.CallbackExecutor;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.os.CancellationSignal;
 import android.os.Looper;
 import android.system.ErrnoException;
 import android.util.Log;
@@ -191,11 +193,18 @@ public final class DnsResolver {
      * @param query blob message
      * @param flags flags as a combination of the FLAGS_* constants
      * @param executor The {@link Executor} that the callback should be executed on.
+     * @param cancellationSignal used by the caller to signal if the query should be
+     *    cancelled. May be {@code null}.
      * @param callback an {@link AnswerCallback} which will be called to notify the caller
      *    of the result of dns query.
      */
     public <T> void query(@Nullable Network network, @NonNull byte[] query, @QueryFlag int flags,
-            @NonNull @CallbackExecutor Executor executor, @NonNull AnswerCallback<T> callback) {
+            @NonNull @CallbackExecutor Executor executor,
+            @Nullable CancellationSignal cancellationSignal,
+            @NonNull AnswerCallback<T> callback) {
+        if (cancellationSignal != null && cancellationSignal.isCanceled()) {
+            return;
+        }
         final FileDescriptor queryfd;
         try {
             queryfd = resNetworkSend((network != null
@@ -205,6 +214,7 @@ public final class DnsResolver {
             return;
         }
 
+        maybeAddCancellationSignal(cancellationSignal, queryfd);
         registerFDListener(executor, queryfd, callback);
     }
 
@@ -219,12 +229,19 @@ public final class DnsResolver {
      * @param nsType dns resource record (RR) type as one of the TYPE_* constants
      * @param flags flags as a combination of the FLAGS_* constants
      * @param executor The {@link Executor} that the callback should be executed on.
+     * @param cancellationSignal used by the caller to signal if the query should be
+     *    cancelled. May be {@code null}.
      * @param callback an {@link AnswerCallback} which will be called to notify the caller
      *    of the result of dns query.
      */
     public <T> void query(@Nullable Network network, @NonNull String domain,
             @QueryClass int nsClass, @QueryType int nsType, @QueryFlag int flags,
-            @NonNull @CallbackExecutor Executor executor, @NonNull AnswerCallback<T> callback) {
+            @NonNull @CallbackExecutor Executor executor,
+            @Nullable CancellationSignal cancellationSignal,
+            @NonNull AnswerCallback<T> callback) {
+        if (cancellationSignal != null && cancellationSignal.isCanceled()) {
+            return;
+        }
         final FileDescriptor queryfd;
         try {
             queryfd = resNetworkQuery((network != null
@@ -233,6 +250,8 @@ public final class DnsResolver {
             callback.onQueryException(e);
             return;
         }
+
+        maybeAddCancellationSignal(cancellationSignal, queryfd);
         registerFDListener(executor, queryfd, callback);
     }
 
@@ -264,6 +283,17 @@ public final class DnsResolver {
                 });
     }
 
+    private void maybeAddCancellationSignal(@Nullable CancellationSignal cancellationSignal,
+            @NonNull FileDescriptor queryfd) {
+        if (cancellationSignal == null) return;
+        cancellationSignal.setOnCancelListener(
+                () -> {
+                    Looper.getMainLooper().getQueue()
+                            .removeOnFileDescriptorEventListener(queryfd);
+                    resNetworkCancel(queryfd);
+            });
+    }
+
     private static class DnsAddressAnswer extends DnsPacket {
         private static final String TAG = "DnsResolver.DnsAddressAnswer";
         private static final boolean DBG = false;

core/java/android/net/NetworkUtils.java

@@ -171,6 +171,12 @@ public class NetworkUtils {
      */
     public static native byte[] resNetworkResult(FileDescriptor fd) throws ErrnoException;
 
+    /**
+     * DNS resolver series jni method.
+     * Attempts to cancel the in-progress query associated with the {@code fd}.
+     */
+    public static native void resNetworkCancel(FileDescriptor fd);
+
     /**
      * Add an entry into the ARP cache.
      */

core/jni/android_net_NetUtils.cpp

@@ -487,6 +487,11 @@ static jbyteArray android_net_utils_resNetworkResult(JNIEnv *env, jobject thiz,
     return answer;
 }
 
+static void android_net_utils_resNetworkCancel(JNIEnv *env, jobject thiz, jobject javaFd) {
+    int fd = jniGetFDFromFileDescriptor(env, javaFd);
+    resNetworkCancel(fd);
+}
+
 static jobject android_net_utils_getTcpRepairWindow(JNIEnv *env, jobject thiz, jobject javaFd) {
     if (javaFd == NULL) {
         jniThrowNullPointerException(env, NULL);
@@ -546,6 +551,7 @@ static const JNINativeMethod gNetworkUtilMethods[] = {
     { "resNetworkSend", "(I[BII)Ljava/io/FileDescriptor;", (void*) android_net_utils_resNetworkSend },
     { "resNetworkQuery", "(ILjava/lang/String;III)Ljava/io/FileDescriptor;", (void*) android_net_utils_resNetworkQuery },
     { "resNetworkResult", "(Ljava/io/FileDescriptor;)[B", (void*) android_net_utils_resNetworkResult },
+    { "resNetworkCancel", "(Ljava/io/FileDescriptor;)V", (void*) android_net_utils_resNetworkCancel },
 };
 
 int register_android_net_NetworkUtils(JNIEnv* env)