feat: Add ability for nat mapping through function (#1948)

Co-authored-by: maxbronnikov10 <[email protected]>
redis · Jan 29, 2025 · 3a04bee · 3a04bee
1 parent 2f9843d
commit 3a04bee
Show file tree

Hide file tree

Showing 6 changed files with 138 additions and 14 deletions.
diff --git a/README.md b/README.md
@@ -1138,7 +1138,31 @@ const cluster = new Redis.Cluster(
 );
 ```
 
+Or you can specify this parameter through function:
+```javascript
+const cluster = new Redis.Cluster(
+  [
+    {
+      host: "203.0.113.73",
+      port: 30001,
+    },
+  ],
+  {
+    natMap: (key) => {
+      if(key.indexOf('30001')) {
+        return { host: "203.0.113.73", port: 30001 };
+      }
+
+      return null;
+    },
+  }
+);
+```
+
 This option is also useful when the cluster is running inside a Docker container.
+Also it works for Clusters in cloud infrastructure where cluster nodes connected through dedicated subnet.
+
+Specifying through may be useful if you don't know concrete internal host and know only node port.
 
 ### Transaction and Pipeline in Cluster Mode
 

diff --git a/lib/cluster/ClusterOptions.ts b/lib/cluster/ClusterOptions.ts
@@ -19,9 +19,11 @@ export type DNSLookupFunction = (
     family?: number
   ) => void
 ) => void;
-export interface NatMap {
+
+export type NatMapFunction = (key: string) => { host: string; port: number } | null;
+export type NatMap = {
   [key: string]: { host: string; port: number };
-}
+} | NatMapFunction
 
 /**
  * Options for Cluster constructor

diff --git a/lib/cluster/index.ts b/lib/cluster/index.ts
@@ -791,17 +791,23 @@ class Cluster extends Commander {
   }
 
   private natMapper(nodeKey: NodeKey | RedisOptions): RedisOptions {
-    if (this.options.natMap && typeof this.options.natMap === "object") {
-      const key =
-        typeof nodeKey === "string"
-          ? nodeKey
-          : `${nodeKey.host}:${nodeKey.port}`;
-      const mapped = this.options.natMap[key];
-      if (mapped) {
-        debug("NAT mapping %s -> %O", key, mapped);
-        return Object.assign({}, mapped);
-      }
+    const key =
+      typeof nodeKey === "string"
+        ? nodeKey
+        : `${nodeKey.host}:${nodeKey.port}`;
+
+    let mapped = null;
+    if (this.options.natMap && typeof this.options.natMap === "function") {
+      mapped = this.options.natMap(key);
+    } else if (this.options.natMap && typeof this.options.natMap === "object") {
+      mapped = this.options.natMap[key];
+    }
+
+    if (mapped) {
+      debug("NAT mapping %s -> %O", key, mapped);
+      return Object.assign({}, mapped);
     }
+
     return typeof nodeKey === "string"
       ? nodeKeyToRedisOptions(nodeKey)
       : nodeKey;

diff --git a/lib/connectors/SentinelConnector/index.ts b/lib/connectors/SentinelConnector/index.ts
@@ -282,7 +282,16 @@ export default class SentinelConnector extends AbstractConnector {
   private sentinelNatResolve(item: SentinelAddress | null) {
     if (!item || !this.options.natMap) return item;
 
-    return this.options.natMap[`${item.host}:${item.port}`] || item;
+    const key = `${item.host}:${item.port}`;
+
+    let result = item;
+    if(typeof this.options.natMap === "function") {
+      result = this.options.natMap(key) || item;
+    } else if (typeof this.options.natMap === "object") {
+      result = this.options.natMap[key] || item;
+    }
+
+    return result;
   }
 
   private connectToSentinel(

diff --git a/test/functional/cluster/nat.ts b/test/functional/cluster/nat.ts
@@ -5,7 +5,7 @@ import { Cluster } from "../../../lib";
 import * as sinon from "sinon";
 
 describe("NAT", () => {
-  it("works for normal case", (done) => {
+  it("works for normal case with object", (done) => {
     const slotTable = [
       [0, 1, ["192.168.1.1", 30001]],
       [2, 16383, ["192.168.1.2", 30001]],
@@ -42,6 +42,48 @@ describe("NAT", () => {
     cluster.get("foo");
   });
 
+  it("works for normal case with function", (done) => {
+    const slotTable = [
+      [0, 1, ["192.168.1.1", 30001]],
+      [2, 16383, ["192.168.1.2", 30001]],
+    ];
+
+    let cluster;
+    new MockServer(30001, null, slotTable);
+    new MockServer(
+      30002,
+      ([command, arg]) => {
+        if (command === "get" && arg === "foo") {
+          cluster.disconnect();
+          done();
+        }
+      },
+      slotTable
+    );
+
+    cluster = new Cluster(
+      [
+        {
+          host: "127.0.0.1",
+          port: 30001,
+        },
+      ],
+      {
+        natMap: (key) => {
+          if(key === "192.168.1.1:30001") {
+            return { host: "127.0.0.1", port: 30001 };
+          }
+          if(key === "192.168.1.2:30001") {
+            return  { host: "127.0.0.1", port: 30002 };
+          }
+          return null;
+        }
+      }
+    );
+
+    cluster.get("foo");
+  });
+
   it("works if natMap does not match all the cases", (done) => {
     const slotTable = [
       [0, 1, ["192.168.1.1", 30001]],

diff --git a/test/unit/clusters/index.ts b/test/unit/clusters/index.ts
@@ -56,6 +56,47 @@ describe("cluster", () => {
       }).to.throw(/Invalid role/);
     });
   });
+
+
+  describe("natMapper", () => {
+    it("returns the original nodeKey if no NAT mapping is provided", () => {
+      const cluster = new Cluster([]);
+      const nodeKey = { host: "127.0.0.1", port: 6379 };
+      const result = cluster["natMapper"](nodeKey);
+
+      expect(result).to.eql(nodeKey);
+    });
+
+    it("maps external IP to internal IP using NAT mapping object", () => {
+      const natMap = { "203.0.113.1:6379": { host: "127.0.0.1", port: 30000 } };
+      const cluster = new Cluster([], { natMap });
+      const nodeKey = "203.0.113.1:6379";
+      const result = cluster["natMapper"](nodeKey);
+      expect(result).to.eql({ host: "127.0.0.1", port: 30000 });
+    });
+
+    it("maps external IP to internal IP using NAT mapping function", () => {
+      const natMap = (key) => {
+        if (key === "203.0.113.1:6379") {
+          return { host: "127.0.0.1", port: 30000 };
+        }
+        return null;
+      };
+      const cluster = new Cluster([], { natMap });
+      const nodeKey = "203.0.113.1:6379";
+      const result = cluster["natMapper"](nodeKey);
+      expect(result).to.eql({ host: "127.0.0.1", port: 30000 });
+    });
+
+    it("returns the original nodeKey if NAT mapping is invalid", () => {
+      const natMap = { "invalid:key": { host: "127.0.0.1", port: 30000 } };
+      const cluster = new Cluster([], { natMap });
+      const nodeKey = "203.0.113.1:6379";
+      const result = cluster["natMapper"](nodeKey);
+      expect(result).to.eql({ host: "203.0.113.1", port: 6379 });
+    });
+  });
+
 });
 
 describe("nodeKeyToRedisOptions()", () => {