From 3b8d00e046d9af93cbc0b269e5d8993b60444c1d Mon Sep 17 00:00:00 2001
From: Ivan Kozlovic <ivan@synadia.com>
Date: Mon, 19 Oct 2020 10:03:00 -0600
Subject: [PATCH] [FIXED] Possible panic when server accepts TLS leafnode
 connection

Signed-off-by: Ivan Kozlovic <ivan@synadia.com>
---
 server/leafnode.go    | 34 +++++++++++++++++++++++++++++++---
 test/leafnode_test.go | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/server/leafnode.go b/server/leafnode.go
index f7d7cd84..c1cb8935 100644
--- a/server/leafnode.go
+++ b/server/leafnode.go
@@ -649,6 +649,14 @@ func (s *Server) createLeafNode(conn net.Conn, remote *leafNodeCfg) *client {
 			return nil
 		}
 
+		// Not sure that can happen, but in case the connection was marked
+		// as closed during the call to parse...
+		if c.isClosed() {
+			c.mu.Unlock()
+			c.closeConnection(ReadError)
+			return nil
+		}
+
 		// Do TLS here as needed.
 		tlsRequired := remote.TLS || remote.TLSConfig != nil
 		if tlsRequired {
@@ -712,6 +720,12 @@ func (s *Server) createLeafNode(conn net.Conn, remote *leafNodeCfg) *client {
 
 			// Re-Grab lock
 			c.mu.Lock()
+
+			// Timeout may have closed the connection while the lock was released.
+			if c.isClosed() {
+				c.mu.Unlock()
+				return nil
+			}
 		}
 
 		if err := c.sendLeafConnect(clusterName, tlsRequired); err != nil {
@@ -736,6 +750,16 @@ func (s *Server) createLeafNode(conn net.Conn, remote *leafNodeCfg) *client {
 		// this before it can initiate the TLS handshake..
 		c.sendProtoNow(bytes.Join(pcs, []byte(" ")))
 
+		// The above call could have marked the connection as closed (due to
+		// TCP error), so if that is the case, bail out here.
+		if c.isClosed() {
+			c.mu.Unlock()
+			// We need to call closeConnection() for proper cleanup, but
+			// "reason" does not really matter since it has been already set.
+			c.closeConnection(WriteError)
+			return nil
+		}
+
 		// Check to see if we need to spin up TLS.
 		if info.TLSRequired {
 			c.Debugf("Starting TLS leafnode server handshake")
@@ -762,14 +786,18 @@ func (s *Server) createLeafNode(conn net.Conn, remote *leafNodeCfg) *client {
 
 			// Indicate that handshake is complete (used in monitoring)
 			c.flags.set(handshakeComplete)
+
+			// Timeout may have closed the connection while the lock was released.
+			if c.isClosed() {
+				c.mu.Unlock()
+				return nil
+			}
 		}
 
 		// Leaf nodes will always require a CONNECT to let us know
 		// when we are properly bound to an account.
 		// The connection may have been closed
-		if !c.isClosed() {
-			c.setAuthTimer(secondsToDuration(opts.LeafNode.AuthTimeout))
-		}
+		c.setAuthTimer(secondsToDuration(opts.LeafNode.AuthTimeout))
 	}
 
 	// Keep track in case server is shutdown before we can successfully register.
diff --git a/test/leafnode_test.go b/test/leafnode_test.go
index c6833713..a8055d8d 100644
--- a/test/leafnode_test.go
+++ b/test/leafnode_test.go
@@ -1204,6 +1204,39 @@ func TestLeafNodeTLS(t *testing.T) {
 	checkLeafNodeConnected(t, s)
 }
 
+func TestLeafNodeTLSConnCloseEarly(t *testing.T) {
+	content := `
+	listen: "127.0.0.1:-1"
+
+	leafnodes {
+		listen: "127.0.0.1:-1"
+		tls {
+			cert_file: "./configs/certs/server-cert.pem"
+			key_file: "./configs/certs/server-key.pem"
+			timeout: 2.0
+		}
+	}
+	`
+	conf := createConfFile(t, []byte(content))
+	defer os.Remove(conf)
+
+	s, opts := RunServerWithConfig(conf)
+	defer s.Shutdown()
+
+	lc, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", opts.LeafNode.Port))
+	if err != nil {
+		t.Fatalf("Unable to connect: %v", err)
+	}
+	// Then close right away
+	lc.Close()
+
+	// Check server does not crash...
+	time.Sleep(250 * time.Millisecond)
+	if s.ID() == "" {
+		t.Fatalf("should not happen")
+	}
+}
+
 type captureLeafNodeErrLogger struct {
 	dummyLogger
 	ch chan string