From 1c5f0330a4a9c581738f77bbbca1dd8843f52b03 Mon Sep 17 00:00:00 2001
From: Nitzan Uziely <nitzan@testim.io>
Date: Fri, 22 Jan 2021 23:01:43 +0200
Subject: [PATCH] dgram: support AbortSignal in createSocket

---
 doc/api/dgram.md                         | 18 +++++++++++++
 lib/dgram.js                             | 15 +++++++++++
 test/parallel/test-dgram-close-signal.js | 33 ++++++++++++++++++++++++
 3 files changed, 66 insertions(+)
 create mode 100644 test/parallel/test-dgram-close-signal.js

diff --git a/doc/api/dgram.md b/doc/api/dgram.md
index d3ed6a6f171621..8dde4670e9c43f 100644
--- a/doc/api/dgram.md
+++ b/doc/api/dgram.md
@@ -735,6 +735,9 @@ chained.
 <!-- YAML
 added: v0.11.13
 changes:
+  - version: REPLACEME
+    pr-url: REPLACEME
+    description: AbortSignal support was added.
   - version: v11.4.0
     pr-url: https://github.com/nodejs/node/pull/23798
     description: The `ipv6Only` option is supported.
@@ -759,6 +762,7 @@ changes:
   * `recvBufferSize` {number} Sets the `SO_RCVBUF` socket value.
   * `sendBufferSize` {number} Sets the `SO_SNDBUF` socket value.
   * `lookup` {Function} Custom lookup function. **Default:** [`dns.lookup()`][].
+  * `signal` {AbortSignal} An AbortSignal that may be used to close a socket.
 * `callback` {Function} Attached as a listener for `'message'` events. Optional.
 * Returns: {dgram.Socket}
 
@@ -770,6 +774,20 @@ method will bind the socket to the "all interfaces" address on a random port
 and port can be retrieved using [`socket.address().address`][] and
 [`socket.address().port`][].
 
+If the `signal` option is enabled, calling `.abort()` on the corresponding
+`AbortController` is similar to calling `.close()` on the socket:
+
+```js
+const controller = new AbortController();
+const { signal } = controller;
+const server = dgram.createSocket({ type: 'udp4', signal });
+server.on('message', (msg, rinfo) => {
+  console.log(`server got: ${msg} from ${rinfo.address}:${rinfo.port}`);
+});
+// Later, when you want to close the server.
+controller.abort();
+```
+
 ### `dgram.createSocket(type[, callback])`
 <!-- YAML
 added: v0.1.99
diff --git a/lib/dgram.js b/lib/dgram.js
index ea1aba522fc938..780a172c4f2249 100644
--- a/lib/dgram.js
+++ b/lib/dgram.js
@@ -52,6 +52,7 @@ const {
 } = errors.codes;
 const {
   isInt32,
+  validateAbortSignal,
   validateString,
   validateNumber,
   validatePort,
@@ -125,6 +126,20 @@ function Socket(type, listener) {
     recvBufferSize,
     sendBufferSize
   };
+
+  if (options?.signal) {
+    const { signal } = options;
+    validateAbortSignal(signal, 'options.signal');
+    const onAborted = () => {
+      this.close();
+    };
+    if (signal.aborted) {
+      process.nextTick(onAborted);
+    } else {
+      signal.addEventListener('abort', onAborted);
+      this.once('close', () => signal.removeEventListener('abort', onAborted));
+    }
+  }
 }
 ObjectSetPrototypeOf(Socket.prototype, EventEmitter.prototype);
 ObjectSetPrototypeOf(Socket, EventEmitter);
diff --git a/test/parallel/test-dgram-close-signal.js b/test/parallel/test-dgram-close-signal.js
new file mode 100644
index 00000000000000..f1008b599dbd18
--- /dev/null
+++ b/test/parallel/test-dgram-close-signal.js
@@ -0,0 +1,33 @@
+'use strict';
+
+const common = require('../common');
+const assert = require('assert');
+const dgram = require('dgram');
+
+{
+  // Test bad signal.
+  assert.throws(
+    () => dgram.createSocket({ type: 'udp4', signal: {} }),
+    {
+      code: 'ERR_INVALID_ARG_TYPE',
+      name: 'TypeError'
+    });
+}
+
+{
+  // Test close.
+  const controller = new AbortController();
+  const { signal } = controller;
+  const server = dgram.createSocket({ type: 'udp4', signal });
+  server.on('close', common.mustCall());
+  controller.abort();
+}
+
+{
+  // Test close with pre-aborted signal.
+  const controller = new AbortController();
+  controller.abort();
+  const { signal } = controller;
+  const server = dgram.createSocket({ type: 'udp4', signal });
+  server.on('close', common.mustCall());
+}