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

This commit is contained in:
Iwasaki Yudai 2015-08-24 19:35:49 +09:00
parent e613b29cc3
commit 470621f39e
2 changed files with 20 additions and 0 deletions

View File

@ -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)

View File

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