Add `--once` option that accepts only one client and exit

app/app.go

@@ -45,6 +45,7 @@ type Options struct {
 	TLSKey        string
 	TLSKey        string
 	TitleFormat   string
 	TitleFormat   string
 	AutoReconnect int
 	AutoReconnect int
+	Once          bool
 	Command       []string
 	Command       []string
 }
 }
 
 
@@ -119,6 +120,11 @@ func (app *App) Run() error {
 		&assetfs.AssetFS{Asset: Asset, AssetDir: AssetDir, Prefix: "static"},
 		&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()
 	var siteMux = http.NewServeMux()
 	siteMux.Handle(path+"/", http.StripPrefix(path+"/", staticHandler))
 	siteMux.Handle(path+"/", http.StripPrefix(path+"/", staticHandler))
 	siteMux.Handle(path+"/ws", wsHandler)
 	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 {
 func generateRandomString(length int) string {
 	const base = 36
 	const base = 36
 	size := big.NewInt(base)
 	size := big.NewInt(base)

main.go

@@ -80,6 +80,11 @@ func main() {
 			Usage:  "Seconds to automatically reconnect to the server when the connection is closed (default: disabled)",
 			Usage:  "Seconds to automatically reconnect to the server when the connection is closed (default: disabled)",
 			EnvVar: "GOTTY_AUTO_RECONNECT",
 			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) {
 	cmd.Action = func(c *cli.Context) {
 		if len(c.Args()) == 0 {
 		if len(c.Args()) == 0 {
@@ -101,6 +106,7 @@ func main() {
 				c.String("tls-key"),
 				c.String("tls-key"),
 				c.String("title-format"),
 				c.String("title-format"),
 				c.Int("auto-reconnect"),
 				c.Int("auto-reconnect"),
+				c.Bool("once"),
 				c.Args(),
 				c.Args(),
 			},
 			},
 		)
 		)