diff --git a/app/app.go b/app/app.go index 1863b5e..6cd2f97 100644 --- a/app/app.go +++ b/app/app.go @@ -122,7 +122,6 @@ func (app *App) Run() error { if app.options.Once { log.Printf("Once option is provided, accepting only one client") - wsHandler = wrapOnce(wsHandler, app) } var siteMux = http.NewServeMux() @@ -272,14 +271,6 @@ 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/app/client_context.go b/app/client_context.go index c36b948..26fb445 100644 --- a/app/client_context.go +++ b/app/client_context.go @@ -62,6 +62,12 @@ func (context *clientContext) goHandleClient() { }() context.app.server.StartRoutine() + + if context.app.options.Once { + log.Printf("Last client accepted, closing the listener.") + context.app.server.Close() + } + go func() { defer context.app.server.FinishRoutine()