From d7a124baaf2ec9d625220922641ae9a4541f39d3 Mon Sep 17 00:00:00 2001
From: Ivan Kozlovic <ivan@synadia.com>
Date: Sun, 15 Aug 2021 12:30:52 -0600
Subject: [PATCH] [FIXED] LeafNode with "wss://.." url was not always
 initiating TLS

If the remote did not have any TLS configuration, the URL scheme
"wss://" was not used as the indicating that the connection should
be attempted as a TLS connection, causing "invalid websocket connection"
in the server attempting to create the remote leafnode connection.

Signed-off-by: Ivan Kozlovic <ivan@synadia.com>
---
 server/leafnode.go      | 17 ++++++++++----
 server/leafnode_test.go | 49 +++++++++++++++++++++++++++++++++++++++++
 server/websocket.go     |  4 ++++
 3 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/server/leafnode.go b/server/leafnode.go
index 515b4f46544..d70ed9ec236 100644
--- a/server/leafnode.go
+++ b/server/leafnode.go
@@ -406,6 +406,11 @@ func newLeafNodeCfg(remote *RemoteLeafOpts) *leafNodeCfg {
 	for _, u := range cfg.urls {
 		cfg.saveTLSHostname(u)
 		cfg.saveUserPassword(u)
+		// If the url(s) have the "wss://" scheme, and we don't have a TLS
+		// config, mark that we should be using TLS anyway.
+		if !cfg.TLS && isWSSURL(u) {
+			cfg.TLS = true
+		}
 	}
 	return cfg
 }
@@ -1136,10 +1141,14 @@ func (c *client) updateLeafNodeURLs(info *Info) {
 	// We have ensured that if a remote has a WS scheme, then all are.
 	// So check if first is WS, then add WS URLs, otherwise, add non WS ones.
 	if len(cfg.URLs) > 0 && isWSURL(cfg.URLs[0]) {
-		// We use wsSchemePrefix. It does not matter if TLS or not since
-		// the distinction is done when creating the LN connection based
-		// on presence of TLS config, etc..
-		c.doUpdateLNURLs(cfg, wsSchemePrefix, info.WSConnectURLs)
+		// It does not really matter if we use "ws://" or "wss://" here since
+		// we will have already marked that the remote should use TLS anyway.
+		// But use proper scheme for log statements, etc...
+		proto := wsSchemePrefix
+		if cfg.TLS {
+			proto = wsSchemePrefixTLS
+		}
+		c.doUpdateLNURLs(cfg, proto, info.WSConnectURLs)
 		return
 	}
 	c.doUpdateLNURLs(cfg, "nats-leaf", info.LeafNodeURLs)
diff --git a/server/leafnode_test.go b/server/leafnode_test.go
index d312144911f..9a862640b11 100644
--- a/server/leafnode_test.go
+++ b/server/leafnode_test.go
@@ -3172,6 +3172,55 @@ func TestLeafNodeWSNoBufferCorruption(t *testing.T) {
 	checkMsgRcv(sub)
 }
 
+func TestLeafNodeWSRemoteNoTLSBlockWithWSSProto(t *testing.T) {
+	o := testDefaultLeafNodeWSOptions()
+	o.Websocket.NoTLS = false
+	tc := &TLSConfigOpts{
+		CertFile: "../test/configs/certs/server-cert.pem",
+		KeyFile:  "../test/configs/certs/server-key.pem",
+		CaFile:   "../test/configs/certs/ca.pem",
+	}
+	tlsConf, err := GenTLSConfig(tc)
+	if err != nil {
+		t.Fatalf("Error generating TLS config: %v", err)
+	}
+	o.Websocket.TLSConfig = tlsConf
+	s := RunServer(o)
+	defer s.Shutdown()
+
+	// The test will make sure that if the protocol is "wss://", a TLS handshake must
+	// be initiated, regardless of the presence of a TLS config block in config file
+	// or here directly.
+	// A bug was causing the absence of TLS config block to initiate a non TLS connection
+	// even if "wss://" proto was specified, which would lead to "invalid websocket connection"
+	// errors in the log.
+	// With the fix, the connection will fail because the remote will fail to verify
+	// the root CA, but at least, we will make sure that this is not an "invalid websocket connection"
+
+	u, _ := url.Parse(fmt.Sprintf("wss://127.0.0.1:%d/some/path", o.Websocket.Port))
+	lo := DefaultOptions()
+	lo.Cluster.Name = "LN"
+	remote := &RemoteLeafOpts{URLs: []*url.URL{u}}
+	lo.LeafNode.Remotes = []*RemoteLeafOpts{remote}
+	lo.LeafNode.ReconnectInterval = 100 * time.Millisecond
+	ln := RunServer(lo)
+	defer ln.Shutdown()
+
+	l := &captureErrorLogger{errCh: make(chan string, 10)}
+	ln.SetLogger(l, false, false)
+
+	select {
+	case e := <-l.errCh:
+		if strings.Contains(e, "invalid websocket connection") {
+			t.Fatalf("The remote did not try to create a TLS connection: %v", e)
+		}
+		// OK!
+		return
+	case <-time.After(2 * time.Second):
+		t.Fatal("Connection should fail")
+	}
+}
+
 func TestLeafNodeStreamImport(t *testing.T) {
 	o1 := DefaultOptions()
 	o1.LeafNode.Port = -1
diff --git a/server/websocket.go b/server/websocket.go
index ab4ce82228d..4b1177a16a8 100644
--- a/server/websocket.go
+++ b/server/websocket.go
@@ -1384,3 +1384,7 @@ func (c *client) wsCollapsePtoNB() (net.Buffers, int64) {
 func isWSURL(u *url.URL) bool {
 	return strings.HasPrefix(strings.ToLower(u.Scheme), wsSchemePrefix)
 }
+
+func isWSSURL(u *url.URL) bool {
+	return strings.HasPrefix(strings.ToLower(u.Scheme), wsSchemePrefixTLS)
+}