From e00dee297d0b82cd2beebcf8cb2d8d0e367fb553 Mon Sep 17 00:00:00 2001
From: theanarkh <theratliter@gmail.com>
Date: Sat, 30 Nov 2024 03:41:22 +0800
Subject: [PATCH] dgram: support blocklist in udp

---
 lib/dgram.js                          | 36 +++++++++++++++++++++++---
 test/parallel/test-dgram-blocklist.js | 37 +++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 3 deletions(-)
 create mode 100644 test/parallel/test-dgram-blocklist.js

diff --git a/lib/dgram.js b/lib/dgram.js
index 74dd83889990ca..2b9b1ff35b0d00 100644
--- a/lib/dgram.js
+++ b/lib/dgram.js
@@ -55,6 +55,7 @@ const {
   _createSocketHandle,
   newHandle,
 } = require('internal/dgram');
+const { isIP } = require('internal/net');
 const {
   isInt32,
   validateAbortSignal,
@@ -104,6 +105,8 @@ function Socket(type, listener) {
   let lookup;
   let recvBufferSize;
   let sendBufferSize;
+  let receiveBlockList;
+  let sendBlockList;
 
   let options;
   if (type !== null && typeof type === 'object') {
@@ -112,6 +115,14 @@ function Socket(type, listener) {
     lookup = options.lookup;
     recvBufferSize = options.recvBufferSize;
     sendBufferSize = options.sendBufferSize;
+    // TODO: validate the params
+    // https://github.com/nodejs/node/pull/56078
+    if (options.receiveBlockList) {
+      receiveBlockList = options.receiveBlockList;
+    }
+    if (options.sendBlockList) {
+      sendBlockList = options.sendBlockList;
+    }
   }
 
   const handle = newHandle(type, lookup);
@@ -134,6 +145,8 @@ function Socket(type, listener) {
     ipv6Only: options?.ipv6Only,
     recvBufferSize,
     sendBufferSize,
+    receiveBlockList,
+    sendBlockList,
   };
 
   if (options?.signal !== undefined) {
@@ -434,9 +447,14 @@ function doConnect(ex, self, ip, address, port, callback) {
     return;
 
   if (!ex) {
-    const err = state.handle.connect(ip, port);
-    if (err) {
-      ex = new ExceptionWithHostPort(err, 'connect', address, port);
+    if (state.sendBlockList?.check(ip, `ipv${isIP(ip)}`)) {
+      // TODO
+      ex = new Error();
+    } else {
+      const err = state.handle.connect(ip, port);
+      if (err) {
+        ex = new ExceptionWithHostPort(err, 'connect', address, port);
+      }
     }
   }
 
@@ -696,6 +714,14 @@ function doSend(ex, self, ip, list, address, port, callback) {
     return;
   }
 
+  if (port && state.sendBlockList?.check(ip, `ipv${isIP(ip)}`)) {
+    if (callback) {
+      // TODO
+      process.nextTick(callback, new Error());
+    }
+    return;
+  }
+
   const req = new SendWrap();
   req.list = list;  // Keep reference alive.
   req.address = address;
@@ -944,6 +970,10 @@ function onMessage(nread, handle, buf, rinfo) {
   if (nread < 0) {
     return self.emit('error', new ErrnoException(nread, 'recvmsg'));
   }
+  if (self[kStateSymbol]?.receiveBlockList?.check(rinfo.address,
+      rinfo.family.toLocaleLowerCase())) {
+    return;
+  }
   rinfo.size = buf.length; // compatibility
   self.emit('message', buf, rinfo);
 }
diff --git a/test/parallel/test-dgram-blocklist.js b/test/parallel/test-dgram-blocklist.js
new file mode 100644
index 00000000000000..9a79738c0c8f18
--- /dev/null
+++ b/test/parallel/test-dgram-blocklist.js
@@ -0,0 +1,37 @@
+'use strict';
+require('../common');
+const common = require('../common');
+const assert = require('assert');
+const dgram = require('dgram');
+const net = require('net');
+
+const blockList = new net.BlockList();
+blockList.addAddress(common.localhostIPv4);
+
+const connectSocket = dgram.createSocket({ type: 'udp4', sendBlockList: blockList });
+connectSocket.bind(0, common.mustCall(() => {
+  connectSocket.connect(9999, common.localhostIPv4, common.mustCall(err => {
+    assert.ok(err);
+    connectSocket.close();
+  }))
+}));
+
+const sendSocket = dgram.createSocket({ type: 'udp4', sendBlockList: blockList });
+sendSocket.bind(0, common.mustCall(() => {
+  sendSocket.send("hello", 9999, common.localhostIPv4, common.mustCall(err => {
+    assert.ok(err);
+    sendSocket.close();
+  }))
+}));
+
+const receiveSocket = dgram.createSocket({ type: 'udp4', receiveBlockList: blockList });
+receiveSocket.om('message', common.mustNotCall());
+receiveSocket.bind(0, common.mustCall(() => {
+  const client = dgram.createSocket('udp4');
+  client.bind(0, common.mustCall(() => {
+    const addressInfo = receiveSocket.address();
+    client.send(addressInfo.port, addressInfo.address, common.mustCall(() => {
+      client.close();
+    }));
+  }))
+}));