feat: add signal option to Deno.resolveDns (#17384)

Closes #14406

5 files changed, 82 insertions, 5 deletions

diff --git a/cli/tests/testdata/run/resolve_dns.ts b/cli/tests/testdata/run/resolve_dns.ts
index ae6ed70a4..a2d0fd046 100644
--- a/cli/tests/testdata/run/resolve_dns.ts
+++ b/cli/tests/testdata/run/resolve_dns.ts
@@ -69,3 +69,25 @@ try {
 } catch (e) {
   console.log(e.message);
 }
+
+try {
+  const ac = new AbortController();
+  queueMicrotask(() => ac.abort());
+  await Deno.resolveDns("www.example.com", "A", {
+    ...nameServer,
+    signal: ac.signal,
+  });
+} catch (e) {
+  console.log(e.name);
+}
+
+try {
+  const ac = new AbortController();
+  ac.abort();
+  await Deno.resolveDns("www.example.com", "A", {
+    ...nameServer,
+    signal: ac.signal,
+  });
+} catch (e) {
+  console.log(e.name);
+}
diff --git a/cli/tests/testdata/run/resolve_dns.ts.out b/cli/tests/testdata/run/resolve_dns.ts.out
index f41dc68bc..025028395 100644
--- a/cli/tests/testdata/run/resolve_dns.ts.out
+++ b/cli/tests/testdata/run/resolve_dns.ts.out
@@ -24,3 +24,5 @@ TXT
 [["I","am","a","txt","record"],["I","am","another","txt","record"],["I am a different","txt record"],["key=val"]]
 Error NotFound thrown for not-found-example.com
 Provided record type is not supported
+AbortError
+AbortError
diff --git a/cli/tsc/dts/lib.deno.ns.d.ts b/cli/tsc/dts/lib.deno.ns.d.ts
index 58dd602d8..33a568b65 100644
--- a/cli/tsc/dts/lib.deno.ns.d.ts
+++ b/cli/tsc/dts/lib.deno.ns.d.ts
@@ -4989,6 +4989,12 @@ declare namespace Deno {
        * @default {53} */
       port?: number;
     };
+    /**
+     * An abort signal to allow cancellation of the DNS resolution operation.
+     * If the signal becomes aborted the resolveDns operation will be stopped
+     * and the promise returned will be rejected with an AbortError.
+     */
+    signal?: AbortSignal;
   }
 
   /** If {@linkcode Deno.resolveDns} is called with `"CAA"` record type
diff --git a/ext/net/01_net.js b/ext/net/01_net.js
index 7b427ef5e..c56f837b8 100644
--- a/ext/net/01_net.js
+++ b/ext/net/01_net.js
@@ -8,6 +8,7 @@ import {
   readableStreamForRidUnrefableUnref,
   writableStreamForRid,
 } from "internal:deno_web/06_streams.js";
+import * as abortSignal from "internal:deno_web/03_abort_signal.js";
 const primordials = globalThis.__bootstrap.primordials;
 const {
   Error,
@@ -30,8 +31,31 @@ function shutdown(rid) {
   return core.shutdown(rid);
 }
 
-function resolveDns(query, recordType, options) {
-  return core.opAsync("op_dns_resolve", { query, recordType, options });
+async function resolveDns(query, recordType, options) {
+  let cancelRid;
+  let abortHandler;
+  if (options?.signal) {
+    options.signal.throwIfAborted();
+    cancelRid = ops.op_cancel_handle();
+    abortHandler = () => core.tryClose(cancelRid);
+    options.signal[abortSignal.add](abortHandler);
+  }
+
+  try {
+    return await core.opAsync("op_dns_resolve", {
+      cancelRid,
+      query,
+      recordType,
+      options,
+    });
+  } finally {
+    if (options?.signal) {
+      options.signal[abortSignal.remove](abortHandler);
+
+      // always throw the abort error when aborted
+      options.signal.throwIfAborted();
+    }
+  }
 }
 
 class Conn {
diff --git a/ext/net/ops.rs b/ext/net/ops.rs
index 1d84a1067..737a46d8c 100644
--- a/ext/net/ops.rs
+++ b/ext/net/ops.rs
@@ -9,6 +9,7 @@ use deno_core::error::custom_error;
 use deno_core::error::generic_error;
 use deno_core::error::AnyError;
 use deno_core::op;
+use deno_core::CancelFuture;
 
 use deno_core::AsyncRefCell;
 use deno_core::ByteString;
@@ -416,6 +417,7 @@ pub enum DnsReturnRecord {
 #[derive(Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct ResolveAddrArgs {
+  cancel_rid: Option<ResourceId>,
   query: String,
   record_type: RecordType,
   options: Option<ResolveDnsOption>,
@@ -451,6 +453,7 @@ where
     query,
     record_type,
     options,
+    cancel_rid,
   } = args;
 
   let (config, opts) = if let Some(name_server) =
@@ -484,9 +487,29 @@ where
 
   let resolver = AsyncResolver::tokio(config, opts)?;
 
-  resolver
-    .lookup(query, record_type)
-    .await
+  let lookup_fut = resolver.lookup(query, record_type);
+
+  let cancel_handle = cancel_rid.and_then(|rid| {
+    state
+      .borrow_mut()
+      .resource_table
+      .get::<CancelHandle>(rid)
+      .ok()
+  });
+
+  let lookup = if let Some(cancel_handle) = cancel_handle {
+    let lookup_rv = lookup_fut.or_cancel(cancel_handle).await;
+
+    if let Some(cancel_rid) = cancel_rid {
+      state.borrow_mut().resource_table.close(cancel_rid).ok();
+    };
+
+    lookup_rv?
+  } else {
+    lookup_fut.await
+  };
+
+  lookup
     .map_err(|e| {
       let message = format!("{e}");
       match e.kind() {