add option for max connection (#112)

add option for max connection

.gotty

@@ -54,6 +54,9 @@
 //       To enable reconnection, set `true` to `enable_reconnect`
 // reconnect_time = false
 
+// [int] Maximum connection to gotty, 0(default) means no limit.
+// max_connection = 0
+
 // [bool] Accept only one client and exit gotty once the client exits
 // once = false
 

app/app.go

@@ -43,6 +43,8 @@ type App struct {
 	titleTemplate *template.Template
 
 	onceMutex *umutex.UnblockingMutex
+
+	connections int
 }
 
 type Options struct {
@@ -62,6 +64,7 @@ type Options struct {
 	TitleFormat         string                 `hcl:"title_format"`
 	EnableReconnect     bool                   `hcl:"enable_reconnect"`
 	ReconnectTime       int                    `hcl:"reconnect_time"`
+	MaxConnection       int                    `hcl:"max_connection"`
 	Once                bool                   `hcl:"once"`
 	PermitArguments     bool                   `hcl:"permit_arguments"`
 	CloseSignal         int                    `hcl:"close_signal"`
@@ -88,6 +91,7 @@ var DefaultOptions = Options{
 	TitleFormat:         "GoTTY - {{ .Command }} ({{ .Hostname }})",
 	EnableReconnect:     false,
 	ReconnectTime:       10,
+	MaxConnection:       0,
 	Once:                false,
 	CloseSignal:         1, // syscall.SIGHUP
 	Preferences:         HtermPrefernces{},
@@ -274,6 +278,12 @@ func (app *App) makeServer(addr string, handler *http.Handler) (*http.Server, er
 }
 
 func (app *App) handleWS(w http.ResponseWriter, r *http.Request) {
+	if app.options.MaxConnection != 0 {
+		if app.connections >= app.options.MaxConnection {
+			log.Printf("Reached max connection: %d", app.options.MaxConnection)
+			return
+		}
+	}
 	log.Printf("New client connected: %s", r.RemoteAddr)
 
 	if r.Method != "GET" {
@@ -342,7 +352,15 @@ func (app *App) handleWS(w http.ResponseWriter, r *http.Request) {
 		log.Print("Failed to execute command")
 		return
 	}
-	log.Printf("Command is running for client %s with PID %d (args=%q)", r.RemoteAddr, cmd.Process.Pid, strings.Join(argv, " "))
+
+	app.connections++
+	if app.options.MaxConnection != 0 {
+		log.Printf("Command is running for client %s with PID %d (args=%q), connections: %d/%d",
+			r.RemoteAddr, cmd.Process.Pid, strings.Join(argv, " "), app.connections, app.options.MaxConnection)
+	} else {
+		log.Printf("Command is running for client %s with PID %d (args=%q), connections: %d",
+			r.RemoteAddr, cmd.Process.Pid, strings.Join(argv, " "), app.connections)
+	}
 
 	context := &clientContext{
 		app:        app,

app/client_context.go

@@ -79,7 +79,14 @@ func (context *clientContext) goHandleClient() {
 
 		context.command.Wait()
 		context.connection.Close()
-		log.Printf("Connection closed: %s", context.request.RemoteAddr)
+		context.app.connections--
+		if context.app.options.MaxConnection != 0 {
+			log.Printf("Connection closed: %s, connections: %d/%d",
+				context.request.RemoteAddr, context.app.connections, context.app.options.MaxConnection)
+		} else {
+			log.Printf("Connection closed: %s, connections: %d",
+				context.request.RemoteAddr, context.app.connections)
+		}
 	}()
 }
 

main.go

@@ -33,6 +33,7 @@ func main() {
 		flag{"title-format", "", "Title format of browser window"},
 		flag{"reconnect", "", "Enable reconnection"},
 		flag{"reconnect-time", "", "Time to reconnect"},
+		flag{"max-connection", "", "Maximum connection to gotty, 0(default) means no limit"},
 		flag{"once", "", "Accept only one client and exit on disconnection"},
 		flag{"permit-arguments", "", "Permit clients to send command line arguments in URL (e.g. http://example.com:8080/?arg=AAA&arg=BBB)"},
 		flag{"close-signal", "", "Signal sent to the command process when gotty close it (default: SIGHUP)"},