server: enforce SSH handshake deadline

Task: 536d2467-2b3d-4b4a-a843-99c96d535cbb

1 files changed, 17 insertions, 0 deletions

diff --git a/internal/server/server.go b/internal/server/server.go
index b4c4406..53aeec6 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"net"
 	"strings"
+	"time"
 
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/dlog"
@@ -18,6 +19,8 @@ import (
 	gossh "golang.org/x/crypto/ssh"
 )
 
+const sshHandshakeTimeout = 10 * time.Second
+
 // Server is the main server data structure.
 type Server struct {
 	cfg config.RuntimeConfig
@@ -118,12 +121,26 @@ func (s *Server) listenerLoop(ctx context.Context, listener net.Listener) {
 func (s *Server) handleConnection(ctx context.Context, conn net.Conn) {
 	dlog.Server.Info("Handling connection")
 
+	// Prevent slow clients from holding connections open indefinitely before SSH handshake completes.
+	if err := conn.SetDeadline(time.Now().Add(sshHandshakeTimeout)); err != nil {
+		dlog.Server.Error("Failed to set SSH handshake deadline", err)
+		conn.Close()
+		return
+	}
+
 	sshConn, chans, reqs, err := gossh.NewServerConn(conn, s.sshServerConfig)
 	if err != nil {
 		dlog.Server.Error("Something just happened", err)
 		return
 	}
 
+	// Handshake succeeded; remove deadline so active sessions are not cut off by the handshake timeout.
+	if err := conn.SetDeadline(time.Time{}); err != nil {
+		dlog.Server.Error("Failed to clear SSH handshake deadline", err)
+		sshConn.Close()
+		return
+	}
+
 	s.stats.incrementConnections()
 	go gossh.DiscardRequests(reqs)
 	for newChannel := range chans {