Add max sessions option to SSH server

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

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")
 

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 {