Close listener after calling `StartRoutine()`

This change prevents the listener from closing itself before
establishing the websocket session with the client.

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)

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