Support query cancellation for async DNS API

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

Change-Id: Iaa72f5c17f58cf0a58663b892bb18cfdf23cd545

2 files changed, 40 insertions, 4 deletions

diff --git a/core/java/android/net/DnsResolver.java b/core/java/android/net/DnsResolver.java
index 93b8cf801d45..59802514c7a3 100644
--- a/core/java/android/net/DnsResolver.java
+++ b/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.
+     *    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.
+     *    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;
diff --git a/core/java/android/net/NetworkUtils.java b/core/java/android/net/NetworkUtils.java
index 0ae29b125149..d2d886ba9a7f 100644
--- a/core/java/android/net/NetworkUtils.java
+++ b/core/java/android/net/NetworkUtils.java
@@ -172,6 +172,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.
      */
     public static void addArpEntry(Inet4Address ipv4Addr, MacAddress ethAddr, String ifname,