diff --git a/bot/bot.go b/bot/bot.go index 28efd32..acfa691 100644 --- a/bot/bot.go +++ b/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 +} diff --git a/bot/session.go b/bot/session.go index 5332e75..8f3b66b 100644 --- a/bot/session.go +++ b/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 } diff --git a/cmd/app.go b/cmd/app.go index 5aeb8c1..1ed42b3 100644 --- a/cmd/app.go +++ b/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 diff --git a/config/config.go b/config/config.go index 9dc7586..126846e 100644 --- a/config/config.go +++ b/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, diff --git a/config/config.yml b/config/config.yml index d181df4..a3f5b5f 100644 --- a/config/config.yml +++ b/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. # diff --git a/config/script.d/helpers/docker-run b/config/script.d/helpers/docker-run index 5b89d04..5138f25 100755 --- a/config/script.d/helpers/docker-run +++ b/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