From 24393f7f61c981f366ccc329be73b4da4e807e91 Mon Sep 17 00:00:00 2001 From: maxbronnikov10 Date: Wed, 22 Jan 2025 00:03:52 +0300 Subject: [PATCH] feat: Add ability for nat mapping through function --- README.md | 24 +++++++++++++ lib/cluster/ClusterOptions.ts | 6 ++-- lib/cluster/index.ts | 26 ++++++++------ lib/connectors/SentinelConnector/index.ts | 11 +++++- test/functional/cluster/nat.ts | 44 ++++++++++++++++++++++- test/unit/clusters/index.ts | 41 +++++++++++++++++++++ 6 files changed, 138 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index 0f61342c..8524d55b 100644 --- a/README.md +++ b/README.md @@ -1132,7 +1132,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 index a829e6e9..26db2e49 100644 --- 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 index c6500f5e..1e76f5e1 100644 --- 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 index be40ef5b..0047a665 100644 --- 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 index 504de35a..10c845b8 100644 --- 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 index 718408e5..f206040b 100644 --- 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()", () => {