From 48a14d3a7570c5d35a7f1ddf674786e4b9cfd5f9 Mon Sep 17 00:00:00 2001
From: abbas ali chezgi <chezgi@gmail.com>
Date: Tue, 20 May 2025 03:18:35 +0330
Subject: [PATCH 1/3] using custom dialer for leaf node connection.

custom dialer is needed when we want to do custom operations like fwmark of connection packets.
---
 server/leafnode.go | 6 +++++-
 server/opts.go     | 7 +++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/server/leafnode.go b/server/leafnode.go
index 4904aee2f21..c580a2b403c 100644
--- a/server/leafnode.go
+++ b/server/leafnode.go
@@ -539,7 +539,11 @@ func (s *Server) connectToRemoteLeafNode(remote *leafNodeCfg, firstConnect bool)
 				err = ErrLeafNodeDisabled
 			} else {
 				s.Debugf("Trying to connect as leafnode to remote server on %q%s", rURL.Host, ipStr)
-				conn, err = natsDialTimeout("tcp", url, dialTimeout)
+				if opts.LeafNode.CustomDialer != nil {
+					conn, err = opts.LeafNode.CustomDialer.Dial("tcp", url)
+				} else {
+					conn, err = natsDialTimeout("tcp", url, dialTimeout)
+				}
 			}
 		}
 		if err != nil {
diff --git a/server/opts.go b/server/opts.go
index f6f7fc18634..d45096e1f53 100644
--- a/server/opts.go
+++ b/server/opts.go
@@ -143,6 +143,10 @@ type RemoteGatewayOpts struct {
 	tlsConfigOpts *TLSConfigOpts
 }
 
+type CustomDialer interface {
+	Dial(network, address string) (net.Conn, error)
+}
+
 // LeafNodeOpts are options for a given server to accept leaf node connections and/or connect to a remote cluster.
 type LeafNodeOpts struct {
 	Host           string        `json:"addr,omitempty"`
@@ -193,6 +197,9 @@ type LeafNodeOpts struct {
 
 	// Snapshot of configured TLS options.
 	tlsConfigOpts *TLSConfigOpts
+
+	// CustomDialer is used to override the default dialer.
+	CustomDialer CustomDialer
 }
 
 // SignatureHandler is used to sign a nonce from the server while

From 3efda025d9680288cb143f07d96b808760f687df Mon Sep 17 00:00:00 2001
From: abbas ali chezgi <chezgi@gmail.com>
Date: Wed, 21 May 2025 05:32:43 +0330
Subject: [PATCH 2/3] adding test for leaf node custom dialer

---
 test/leafnode_test.go | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/test/leafnode_test.go b/test/leafnode_test.go
index 3fb01902bab..0eac2b600c7 100644
--- a/test/leafnode_test.go
+++ b/test/leafnode_test.go
@@ -76,6 +76,18 @@ func runSolicitLeafServerToURL(surl string) (*server.Server, *server.Options) {
 	return RunServer(&o), &o
 }
 
+func runSolicitLeafServerCustomDialer(lso *server.Options,dialer server.CustomDialer) (*server.Server, *server.Options) {
+	surl := fmt.Sprintf("nats-leaf://%s:%d", lso.LeafNode.Host, lso.LeafNode.Port
+	o := DefaultTestOptions
+	o.Port = -1
+	o.NoSystemAccount = true
+	rurl, _ := url.Parse(surl)
+	o.LeafNode.Remotes = []*server.RemoteLeafOpts{{URLs: []*url.URL{rurl}}}
+	o.LeafNode.ReconnectInterval = 100 * time.Millisecond
+	o.LeafNode.CustomDialer = dialer
+	return RunServer(&o), &o
+}
+
 func TestLeafNodeInfo(t *testing.T) {
 	s, opts := runLeafServer()
 	defer s.Shutdown()
@@ -526,6 +538,34 @@ func TestLeafNodeSolicit(t *testing.T) {
 	checkLeafNodeConnected(t, s)
 }
 
+type TestDialer struct{
+}
+
+func (td TestDialer) Dial(network, address string) (net.Conn, error) {
+	d := net.Dialer{
+		Timeout:   1 * time.Second,
+		KeepAlive: -1,
+	}
+	return d.Dial(network, address)
+}
+
+func TestLeafNodeSolicitCustomDialer(t *testing.T) {
+	s, opts := runLeafServer()
+	defer s.Shutdown()
+
+	sl, _ := runSolicitLeafServerCustomDialer(opts,&TestDialer{})
+	defer sl.Shutdown()
+
+	checkLeafNodeConnected(t, s)
+
+	// Now test reconnect.
+	s.Shutdown()
+	// Need to restart it on the same port.
+	s, _ = runLeafServerOnPort(opts.LeafNode.Port)
+	defer s.Shutdown()
+	checkLeafNodeConnected(t, s)
+}
+
 func TestLeafNodeNoEcho(t *testing.T) {
 	s, opts := runLeafServer()
 	defer s.Shutdown()

From 363cd4ffa57d393db203a97c20000cefedc4f770 Mon Sep 17 00:00:00 2001
From: abbas ali chezgi <chezgi@gmail.com>
Date: Wed, 21 May 2025 11:24:52 +0330
Subject: [PATCH 3/3] lint bug fixed

---
 test/leafnode_test.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/leafnode_test.go b/test/leafnode_test.go
index 0eac2b600c7..8530b25e3bf 100644
--- a/test/leafnode_test.go
+++ b/test/leafnode_test.go
@@ -76,8 +76,8 @@ func runSolicitLeafServerToURL(surl string) (*server.Server, *server.Options) {
 	return RunServer(&o), &o
 }
 
-func runSolicitLeafServerCustomDialer(lso *server.Options,dialer server.CustomDialer) (*server.Server, *server.Options) {
-	surl := fmt.Sprintf("nats-leaf://%s:%d", lso.LeafNode.Host, lso.LeafNode.Port
+func runSolicitLeafServerCustomDialer(lso *server.Options, dialer server.CustomDialer) (*server.Server, *server.Options) {
+	surl := fmt.Sprintf("nats-leaf://%s:%d", lso.LeafNode.Host, lso.LeafNode.Port)
 	o := DefaultTestOptions
 	o.Port = -1
 	o.NoSystemAccount = true