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(), }, )