Add --max-total-sessions and --max-user-sessions

bot/bot.go

@@ -28,13 +28,15 @@ const (
 		"recorded (default: `%s`). To start a private REPL session, just DM me."
 	shareMessage = "Using the word `share` will allow you to share your own terminal here in the chat. Terminal sharing " +
 		"sessions are always started in `only-me` mode, unless overridden."
-	unknownCommandMessage = "I am not quite sure what you mean by _%s_ ⁉"
-	misconfiguredMessage  = "😭 Oh no. It looks like REPLbot is misconfigured. I couldn't find any scripts to run."
-	helpRequestedCommand  = "help"
-	recordCommand         = "record"
-	noRecordCommand       = "norecord"
-	shareCommand          = "share"
-	shareServerScriptFile = "/tmp/replbot_share_server.sh"
+	unknownCommandMessage           = "I am not quite sure what you mean by _%s_ ⁉"
+	misconfiguredMessage            = "😭 Oh no. It looks like REPLbot is misconfigured. I couldn't find any scripts to run."
+	maxTotalSessionsExceededMessage = "😭 There are too many active sessions. Please wait until another session is closed."
+	maxUserSessionsExceededMessage  = "😭 You have too many active sessions. Please close a session to start a new one."
+	helpRequestedCommand            = "help"
+	recordCommand                   = "record"
+	noRecordCommand                 = "norecord"
+	shareCommand                    = "share"
+	shareServerScriptFile           = "/tmp/replbot_share_server.sh"
 )
 
 // Key exchange algorithms, ciphers,and MACs (see `ssh-audit` output)
@@ -164,6 +166,9 @@ func (b *Bot) handleMessageEvent(ev *messageEvent) error {
 	if err != nil {
 		return b.handleHelp(ev.Channel, ev.Thread, err)
 	}
+	if allowed, err := b.checkSessionAllowed(ev.Channel, ev.Thread, conf); err != nil || !allowed {
+		return err
+	}
 	switch conf.controlMode {
 	case config.Channel:
 		return b.startSessionChannel(ev, conf)
@@ -479,3 +484,23 @@ func (b *Bot) sshServerConfigCallback(ctx ssh.Context) *gossh.ServerConfig {
 	conf.MACs = []string{macHMACSHA256ETM}
 	return conf
 }
+
+func (b *Bot) checkSessionAllowed(channel, thread string, conf *sessionConfig) (allowed bool, err error) {
+	b.mu.RLock()
+	defer b.mu.RUnlock()
+	if len(b.sessions) >= b.config.MaxTotalSessions {
+		ch := &channelID{Channel: channel, Thread: thread}
+		return false, b.conn.Send(ch, maxTotalSessionsExceededMessage)
+	}
+	var userSessions int
+	for _, sess := range b.sessions {
+		if sess.conf.user == conf.user {
+			userSessions++
+		}
+	}
+	if userSessions >= b.config.MaxUserSessions {
+		ch := &channelID{Channel: channel, Thread: thread}
+		return false, b.conn.Send(ch, maxUserSessionsExceededMessage)
+	}
+	return true, nil
+}

bot/session.go

@@ -524,9 +524,10 @@ func (s *session) getEnv() (map[string]string, error) {
 		}
 	}
 	return map[string]string{
-		"REPLBOT_SSH_KEY_FILE":   sshKeyFile,
-		"REPLBOT_SSH_USER_FILE":  sshUserFile,
-		"REPLBOT_SSH_RELAY_PORT": relayPort,
+		"REPLBOT_SSH_KEY_FILE":       sshKeyFile,
+		"REPLBOT_SSH_USER_FILE":      sshUserFile,
+		"REPLBOT_SSH_RELAY_PORT":     relayPort,
+		"REPLBOT_MAX_TOTAL_SESSIONS": strconv.Itoa(s.conf.global.MaxUserSessions),
 	}, nil
 }
 

cmd/app.go

@@ -24,6 +24,8 @@ func New() *cli.App {
 		altsrc.NewStringFlag(&cli.StringFlag{Name: "bot-token", Aliases: []string{"t"}, EnvVars: []string{"REPLBOT_BOT_TOKEN"}, DefaultText: "none", Usage: "bot token"}),
 		altsrc.NewStringFlag(&cli.StringFlag{Name: "script-dir", Aliases: []string{"d"}, EnvVars: []string{"REPLBOT_SCRIPT_DIR"}, Value: "/etc/replbot/script.d", DefaultText: "/etc/replbot/script.d", Usage: "script directory"}),
 		altsrc.NewDurationFlag(&cli.DurationFlag{Name: "idle-timeout", Aliases: []string{"T"}, EnvVars: []string{"REPLBOT_IDLE_TIMEOUT"}, Value: config.DefaultIdleTimeout, Usage: "timeout after which sessions are ended"}),
+		altsrc.NewIntFlag(&cli.IntFlag{Name: "max-total-sessions", Aliases: []string{"S"}, EnvVars: []string{"REPLBOT_MAX_TOTAL_SESSIONS"}, Value: config.DefaultMaxTotalSessions, Usage: "max number of concurrent total sessions"}),
+		altsrc.NewIntFlag(&cli.IntFlag{Name: "max-user-sessions", Aliases: []string{"U"}, EnvVars: []string{"REPLBOT_MAX_USER_SESSIONS"}, Value: config.DefaultMaxUserSessions, Usage: "max number of concurrent sessions per user"}),
 		altsrc.NewStringFlag(&cli.StringFlag{Name: "default-control-mode", Aliases: []string{"m"}, EnvVars: []string{"REPLBOT_DEFAULT_CONTROL_MODE"}, Value: string(config.DefaultControlMode), DefaultText: string(config.DefaultControlMode), Usage: "default control mode [channel, thread or split]"}),
 		altsrc.NewStringFlag(&cli.StringFlag{Name: "default-window-mode", Aliases: []string{"w"}, EnvVars: []string{"REPLBOT_DEFAULT_WINDOW_MODE"}, Value: string(config.DefaultWindowMode), DefaultText: string(config.DefaultWindowMode), Usage: "default window mode [full or trim]"}),
 		altsrc.NewStringFlag(&cli.StringFlag{Name: "default-auth-mode", Aliases: []string{"a"}, EnvVars: []string{"REPLBOT_DEFAULT_AUTH_MODE"}, Value: string(config.DefaultAuthMode), DefaultText: string(config.DefaultAuthMode), Usage: "default auth mode [only-me or everyone]"}),
@@ -55,6 +57,8 @@ func execRun(c *cli.Context) error {
 	token := c.String("bot-token")
 	scriptDir := c.String("script-dir")
 	timeout := c.Duration("idle-timeout")
+	maxTotalSessions := c.Int("max-total-sessions")
+	maxUserSessions := c.Int("max-user-sessions")
 	defaultControlMode := config.ControlMode(c.String("default-control-mode"))
 	defaultWindowMode := config.WindowMode(c.String("default-window-mode"))
 	defaultAuthMode := config.AuthMode(c.String("default-auth-mode"))
@@ -78,6 +82,8 @@ func execRun(c *cli.Context) error {
 		return errors.New("default window mode must be 'full' or 'trim'")
 	} else if shareHost != "" && (shareKeyFile == "" || !util.FileExists(shareKeyFile)) {
 		return errors.New("share key file must be set and exist if share host is set, check --share-key-file or REPLBOT_SHARE_KEY_FILE")
+	} else if maxUserSessions > maxTotalSessions {
+		return errors.New("max total sessions must be larger or equal to max user sessions")
 	}
 	cursorRate, err := parseCursorRate(cursor)
 	if err != nil {
@@ -100,6 +106,8 @@ func execRun(c *cli.Context) error {
 	conf := config.New(token)
 	conf.ScriptDir = scriptDir
 	conf.IdleTimeout = timeout
+	conf.MaxTotalSessions = maxTotalSessions
+	conf.MaxUserSessions = maxUserSessions
 	conf.DefaultControlMode = defaultControlMode
 	conf.DefaultWindowMode = defaultWindowMode
 	conf.DefaultAuthMode = defaultAuthMode

config/config.go

@@ -12,6 +12,12 @@ const (
 	// DefaultIdleTimeout defines the default time after which a session is terminated
 	DefaultIdleTimeout = 10 * time.Minute
 
+	// DefaultMaxTotalSessions is the default number of sessions all users are allowed to run concurrently
+	DefaultMaxTotalSessions = 6
+
+	// DefaultMaxUserSessions is the default number of sessions a user is allowed to run concurrently
+	DefaultMaxUserSessions = 2
+
 	// DefaultRecord defines if sessions are recorded by default
 	DefaultRecord = false
 
@@ -24,6 +30,8 @@ type Config struct {
 	Token              string
 	ScriptDir          string
 	IdleTimeout        time.Duration
+	MaxTotalSessions   int
+	MaxUserSessions    int
 	DefaultControlMode ControlMode
 	DefaultWindowMode  WindowMode
 	DefaultAuthMode    AuthMode
@@ -41,6 +49,8 @@ func New(token string) *Config {
 	return &Config{
 		Token:              token,
 		IdleTimeout:        DefaultIdleTimeout,
+		MaxTotalSessions:   DefaultMaxTotalSessions,
+		MaxUserSessions:    DefaultMaxUserSessions,
 		DefaultControlMode: DefaultControlMode,
 		DefaultWindowMode:  DefaultWindowMode,
 		DefaultAuthMode:    DefaultAuthMode,

config/config.yml

@@ -102,6 +102,9 @@ bot-token: MUST_BE_SET
 #
 # idle-timeout: 10m
 
+# max-total-sessions: 6
+# max-user-sessions: 2
+
 # Cursor setting for the terminal. Can be "on" to always render the cursor, "off" to turn it off entirely,
 # or a duration such as "1s" or "2s" to define the blink rate.
 #

config/script.d/helpers/docker-run

@@ -16,9 +16,12 @@ container="$2"
 shift; shift
 
 # Determine CPU/RAM resources
-divisor=4
+if [ -n "${REPLBOT_MAX_TOTAL_SESSIONS}" ]; then
+  divisor="${REPLBOT_MAX_TOTAL_SESSIONS}"
+else
+  divisor=4
+fi
 cpus="$(grep -c '^processor' /proc/cpuinfo)"
-mem="$(grep MemTotal /proc/meminfo | awk '{print $2}')"
 if [ -z "${cpus}" ]; then
   cpus=0.5
 else
@@ -27,6 +30,7 @@ else
     cpus=0.5
   fi
 fi
+mem="$(grep MemTotal /proc/meminfo | awk '{print $2}')"
 if [ -z "${mem}" ]; then
   mem=128
 else