From 470621f39ed98fc20f225126acb14cfa979c82a3 Mon Sep 17 00:00:00 2001
From: Iwasaki Yudai <yudai@arielworks.com>
Date: Mon, 24 Aug 2015 19:35:49 +0900
Subject: [PATCH] Add `--once` option that accepts only one client and exit

---
 app/app.go | 14 ++++++++++++++
 main.go    |  6 ++++++
 2 files changed, 20 insertions(+)

diff --git a/app/app.go b/app/app.go
index 84f3d11..1863b5e 100644
--- a/app/app.go
+++ b/app/app.go
@@ -45,6 +45,7 @@ type Options struct {
 	TLSKey        string
 	TitleFormat   string
 	AutoReconnect int
+	Once          bool
 	Command       []string
 }
 
@@ -119,6 +120,11 @@ func (app *App) Run() error {
 		&assetfs.AssetFS{Asset: Asset, AssetDir: AssetDir, Prefix: "static"},
 	)
 
+	if app.options.Once {
+		log.Printf("Once option is provided, accepting only one client")
+		wsHandler = wrapOnce(wsHandler, app)
+	}
+
 	var siteMux = http.NewServeMux()
 	siteMux.Handle(path+"/", http.StripPrefix(path+"/", staticHandler))
 	siteMux.Handle(path+"/ws", wsHandler)
@@ -266,6 +272,14 @@ func wrapBasicAuth(handler http.Handler, credential string) http.Handler {
 	})
 }
 
+func wrapOnce(handler http.HandlerFunc, app *App) http.HandlerFunc {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		log.Printf("Last client accepted, closing the listener.")
+		app.server.Close()
+		handler.ServeHTTP(w, r)
+	})
+}
+
 func generateRandomString(length int) string {
 	const base = 36
 	size := big.NewInt(base)
diff --git a/main.go b/main.go
index 88d9ce1..ac9e347 100644
--- a/main.go
+++ b/main.go
@@ -80,6 +80,11 @@ func main() {
 			Usage:  "Seconds to automatically reconnect to the server when the connection is closed (default: disabled)",
 			EnvVar: "GOTTY_AUTO_RECONNECT",
 		},
+		cli.BoolFlag{
+			Name:   "once",
+			Usage:  "Accept only one client and exits on disconnection",
+			EnvVar: "GOTTY_ONCE",
+		},
 	}
 	cmd.Action = func(c *cli.Context) {
 		if len(c.Args()) == 0 {
@@ -101,6 +106,7 @@ func main() {
 				c.String("tls-key"),
 				c.String("title-format"),
 				c.Int("auto-reconnect"),
+				c.Bool("once"),
 				c.Args(),
 			},
 		)