SSH handshake in goroutine and defer close (#1640)

* SSH handshake in goroutine and defer close

sshd/server.go

@@ -2,10 +2,10 @@ package sshd
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"net"
-	"sync"
 
 	"github.com/armon/go-radix"
 	"github.com/sirupsen/logrus"
@@ -27,20 +27,21 @@ type SSHServer struct {
 	commands    *radix.Tree
 	listener    net.Listener
 
-	// Locks the conns/counter to avoid concurrent map access
-	connsLock sync.Mutex
-	conns     map[int]*session
-	counter   int
+	// Call the cancel() function to stop all active sessions
+	ctx    context.Context
+	cancel func()
 }
 
 // NewSSHServer creates a new ssh server rigged with default commands and prepares to listen
 func NewSSHServer(l *logrus.Entry) (*SSHServer, error) {
 
+	ctx, cancel := context.WithCancel(context.Background())
 	s := &SSHServer{
 		trustedKeys: make(map[string]map[string]bool),
 		l:           l,
 		commands:    radix.New(),
-		conns:       make(map[int]*session),
+		ctx:         ctx,
+		cancel:      cancel,
 	}
 
 	cc := ssh.CertChecker{
@@ -175,44 +176,44 @@ func (s *SSHServer) run() {
 			}
 			return
 		}
-
-		conn, chans, reqs, err := ssh.NewServerConn(c, s.config)
-		fp := ""
-		if conn != nil {
-			fp = conn.Permissions.Extensions["fp"]
-		}
-
-		if err != nil {
-			l := s.l.WithError(err).WithField("remoteAddress", c.RemoteAddr())
+		go func(c net.Conn) {
+			// NewServerConn may block while waiting for the client to complete the handshake.
+			// Ensure that a bad client doesn't hurt us by checking for the parent context
+			// cancellation before calling NewServerConn, and forcing the socket to close when
+			// the context is cancelled.
+			sessionContext, sessionCancel := context.WithCancel(s.ctx)
+			go func() {
+				<-sessionContext.Done()
+				c.Close()
+			}()
+			conn, chans, reqs, err := ssh.NewServerConn(c, s.config)
+			fp := ""
 			if conn != nil {
-				l = l.WithField("sshUser", conn.User())
-				conn.Close()
+				fp = conn.Permissions.Extensions["fp"]
 			}
-			if fp != "" {
-				l = l.WithField("sshFingerprint", fp)
+
+			if err != nil {
+				l := s.l.WithError(err).WithField("remoteAddress", c.RemoteAddr())
+				if conn != nil {
+					l = l.WithField("sshUser", conn.User())
+					conn.Close()
+				}
+				if fp != "" {
+					l = l.WithField("sshFingerprint", fp)
+				}
+				l.Warn("failed to handshake")
+				sessionCancel()
+				return
 			}
-			l.Warn("failed to handshake")
-			continue
-		}
 
-		l := s.l.WithField("sshUser", conn.User())
-		l.WithField("remoteAddress", c.RemoteAddr()).WithField("sshFingerprint", fp).Info("ssh user logged in")
+			l := s.l.WithField("sshUser", conn.User())
+			l.WithField("remoteAddress", c.RemoteAddr()).WithField("sshFingerprint", fp).Info("ssh user logged in")
 
-		session := NewSession(s.commands, conn, chans, l.WithField("subsystem", "sshd.session"))
-		s.connsLock.Lock()
-		s.counter++
-		counter := s.counter
-		s.conns[counter] = session
-		s.connsLock.Unlock()
+			NewSession(s.commands, conn, chans, sessionCancel, l.WithField("subsystem", "sshd.session"))
 
-		go ssh.DiscardRequests(reqs)
-		go func() {
-			<-session.exitChan
-			s.l.WithField("id", counter).Debug("closing conn")
-			s.connsLock.Lock()
-			delete(s.conns, counter)
-			s.connsLock.Unlock()
-		}()
+			go ssh.DiscardRequests(reqs)
+
+		}(c)
 	}
 }
 
@@ -226,9 +227,5 @@ func (s *SSHServer) Stop() {
 }
 
 func (s *SSHServer) closeSessions() {
-	s.connsLock.Lock()
-	for _, c := range s.conns {
-		c.Close()
-	}
-	s.connsLock.Unlock()
+	s.cancel()
 }