From be07d420dd91586251e2efbc226e8f055c08392f Mon Sep 17 00:00:00 2001
From: Yifa Zhang <zyfdegh@users.noreply.github.com>
Date: Sat, 13 Aug 2016 15:29:21 +0800
Subject: [PATCH] add option for max connection (#112)

add option for max connection
---
 .gotty                |  3 +++
 app/app.go            | 20 +++++++++++++++++++-
 app/client_context.go |  9 ++++++++-
 main.go               |  1 +
 4 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/.gotty b/.gotty
index 51d6e22..a3c89a2 100644
--- a/.gotty
+++ b/.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
 
diff --git a/app/app.go b/app/app.go
index 60722cc..335c19b 100644
--- a/app/app.go
+++ b/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,
diff --git a/app/client_context.go b/app/client_context.go
index d4c2f29..3a3b27a 100644
--- a/app/client_context.go
+++ b/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)
+		}
 	}()
 }
 
diff --git a/main.go b/main.go
index 20f3600..7d14198 100644
--- a/main.go
+++ b/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)"},