From 2b6b7b629b6cb50f9265ff9fab11cec6ca1c401e Mon Sep 17 00:00:00 2001
From: Miguel Mota <miguelmota2@gmail.com>
Date: Tue, 16 Mar 2021 23:18:44 -0700
Subject: [PATCH] Add max sessions option to SSH server

---
 Makefile               |  6 ++++++
 cmd/commands/server.go |  3 +++
 pkg/ssh/server.go      | 16 ++++++++++++++++
 3 files changed, 25 insertions(+)

diff --git a/Makefile b/Makefile
index 1400e67..80818b6 100644
--- a/Makefile
+++ b/Makefile
@@ -238,5 +238,11 @@ docker-build-and-push: docker-build docker-push
 docker-run-ssh:
 	docker run -p 2222:22 -v ~/.ssh/demo:/keys --entrypoint cointop -it cointop/cointop server -k /keys/id_rsa
 
+ssh-server:
+	go run cmd/cointop/cointop.go server -p 2222
+
+ssh-client:
+	ssh localhost -p 2222
+
 mp3:
 	cat <(printf "package notifier\nfunc Mp3() string {\nreturn \`" "") <(xxd -p media/notification.mp3 | tr -d "\n") <(printf "\`\n}" "") > pkg/notifier/mp3.go
diff --git a/cmd/commands/server.go b/cmd/commands/server.go
index f1fe526..fbd8ce3 100644
--- a/cmd/commands/server.go
+++ b/cmd/commands/server.go
@@ -16,6 +16,7 @@ func ServerCmd() *cobra.Command {
 	var address string = "0.0.0.0"
 	var idleTimeout uint = 0
 	var maxTimeout uint = 0
+	var maxSessions uint = 0
 	var executableBinary string = "cointop"
 	var hostKeyFile string = cssh.DefaultHostKeyFile
 
@@ -29,6 +30,7 @@ func ServerCmd() *cobra.Command {
 				Port:             port,
 				IdleTimeout:      time.Duration(int(idleTimeout)) * time.Second,
 				MaxTimeout:       time.Duration(int(maxTimeout)) * time.Second,
+				MaxSessions:      maxSessions,
 				ExecutableBinary: executableBinary,
 				HostKeyFile:      hostKeyFile,
 			})
@@ -42,6 +44,7 @@ func ServerCmd() *cobra.Command {
 	serverCmd.Flags().StringVarP(&address, "address", "a", address, "Address")
 	serverCmd.Flags().UintVarP(&idleTimeout, "idle-timeout", "t", idleTimeout, "Idle timeout in seconds. Default is 0 for no idle timeout")
 	serverCmd.Flags().UintVarP(&maxTimeout, "max-timeout", "m", maxTimeout, "Max timeout in seconds. Default is 0 for no max timeout")
+	serverCmd.Flags().UintVarP(&maxSessions, "max-sessions", "", maxSessions, "Max number of sessions allowed. Default is 0 for unlimited.")
 	serverCmd.Flags().StringVarP(&executableBinary, "binary", "b", executableBinary, "Executable binary path")
 	serverCmd.Flags().StringVarP(&hostKeyFile, "host-key-file", "k", hostKeyFile, "Host key file")
 
diff --git a/pkg/ssh/server.go b/pkg/ssh/server.go
index d84949d..42f06a6 100644
--- a/pkg/ssh/server.go
+++ b/pkg/ssh/server.go
@@ -32,6 +32,7 @@ type Config struct {
 	MaxTimeout       time.Duration
 	ExecutableBinary string
 	HostKeyFile      string
+	MaxSessions      uint
 }
 
 // Server is server struct
@@ -43,6 +44,8 @@ type Server struct {
 	executableBinary string
 	sshServer        *ssh.Server
 	hostKeyFile      string
+	maxSessions      uint
+	sessionCount     uint
 }
 
 // NewServer returns a new server instance
@@ -60,6 +63,7 @@ func NewServer(config *Config) *Server {
 		maxTimeout:       config.MaxTimeout,
 		executableBinary: config.ExecutableBinary,
 		hostKeyFile:      hostKeyFile,
+		maxSessions:      config.MaxSessions,
 	}
 }
 
@@ -70,6 +74,18 @@ func (s *Server) ListenAndServe() error {
 		IdleTimeout: s.idleTimeout,
 		MaxTimeout:  s.maxTimeout,
 		Handler: func(sshSession ssh.Session) {
+			if s.maxSessions > 0 {
+				s.sessionCount++
+				defer func() {
+					s.sessionCount--
+				}()
+				if s.sessionCount > s.maxSessions {
+					io.WriteString(sshSession, "Error: Maximum sessions reached. Must wait until session slot is available.")
+					sshSession.Exit(1)
+					return
+				}
+			}
+
 			cmdUserArgs := sshSession.Command()
 			ptyReq, winCh, isPty := sshSession.Pty()
 			if !isPty {