1 files changed, 59 insertions, 4 deletions

diff --git a/internal/server/stats.go b/internal/server/stats.go
index 99a644a..7a60d61 100644
--- a/internal/server/stats.go
+++ b/internal/server/stats.go
@@ -6,7 +6,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/dlog"
 )
 
@@ -15,6 +14,20 @@ type stats struct {
 	mutex               sync.Mutex
 	currentConnections  int
 	lifetimeConnections uint64
+	maxConnections      int
+	// preAuthConnections counts TCP connections that have been accepted but
+	// whose SSH handshake has not yet completed. These are counted against
+	// maxConnections so that slow or abusive clients cannot create unbounded
+	// goroutines during the handshake phase. Once a handshake succeeds the
+	// slot is converted to a regular connection (decrementPreAuth +
+	// incrementConnections). On failure the slot is simply released.
+	preAuthConnections int
+}
+
+func newStats(maxConnections int) stats {
+	return stats{
+		maxConnections: maxConnections,
+	}
 }
 
 func (s *stats) incrementConnections() {
@@ -32,6 +45,40 @@ func (s *stats) decrementConnections() {
 	s.mutex.Unlock()
 }
 
+// reservePreAuth increments the pre-auth counter immediately after Accept so
+// that slow or unauthenticated handshakes are counted against maxConnections.
+// It must be paired with exactly one call to releasePreAuth or
+// promotePreAuthToConnection.
+func (s *stats) reservePreAuth() {
+	defer s.logServerStats()
+	s.mutex.Lock()
+	s.preAuthConnections++
+	s.mutex.Unlock()
+}
+
+// releasePreAuth decrements the pre-auth counter without converting the slot
+// into a full authenticated connection. Call this on every handshake failure
+// path to undo the reservePreAuth reservation.
+func (s *stats) releasePreAuth() {
+	defer s.logServerStats()
+	s.mutex.Lock()
+	s.preAuthConnections--
+	s.mutex.Unlock()
+}
+
+// promotePreAuthToConnection atomically converts a pre-auth reservation into a
+// full authenticated connection. It decrements preAuthConnections and
+// increments both currentConnections and lifetimeConnections under a single
+// lock acquisition so there is no instant where neither counter holds the slot.
+func (s *stats) promotePreAuthToConnection() {
+	defer s.logServerStats()
+	s.mutex.Lock()
+	s.preAuthConnections--
+	s.currentConnections++
+	s.lifetimeConnections++
+	s.mutex.Unlock()
+}
+
 func (s *stats) hasConnections() bool {
 	s.mutex.Lock()
 	currentConnections := s.currentConnections
@@ -50,16 +97,24 @@ func (s *stats) logServerStats() {
 	data := make(map[string]interface{})
 	data["currentConnections"] = s.currentConnections
 	data["lifetimeConnections"] = s.lifetimeConnections
+	data["preAuthConnections"] = s.preAuthConnections
 	dlog.Server.Mapreduce("STATS", data)
 }
 
+// serverLimitExceeded checks whether accepting another connection would exceed
+// maxConnections. Both authenticated connections (currentConnections) and
+// in-progress handshakes (preAuthConnections) are counted so that slow or
+// unauthenticated clients cannot bypass the limit by keeping many TCP
+// connections open during the handshake phase.
 func (s *stats) serverLimitExceeded() error {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
-	if s.currentConnections >= config.Server.MaxConnections {
-		return fmt.Errorf("Exceeded max allowed concurrent connections of %d",
-			config.Server.MaxConnections)
+	// Count both authenticated connections and pre-auth handshakes in progress.
+	total := s.currentConnections + s.preAuthConnections
+	if total >= s.maxConnections {
+		return fmt.Errorf("Exceeded max allowed concurrent connections of %d (current=%d, pre-auth=%d)",
+			s.maxConnections, s.currentConnections, s.preAuthConnections)
 	}
 	return nil
 }