Ensure only one client can connect when `--once` is given

Using a mutex

app/app.go

@@ -25,6 +25,7 @@ import (
 	"github.com/gorilla/websocket"
 	"github.com/gorilla/websocket"
 	"github.com/hashicorp/hcl"
 	"github.com/hashicorp/hcl"
 	"github.com/kr/pty"
 	"github.com/kr/pty"
+	"github.com/yudai/umutex"
 )
 )
 
 
 type InitMessage struct {
 type InitMessage struct {
@@ -40,6 +41,8 @@ type App struct {
 	server   *manners.GracefulServer
 	server   *manners.GracefulServer
 
 
 	titleTemplate *template.Template
 	titleTemplate *template.Template
+
+	onceMutex *umutex.UnblockingMutex
 }
 }
 
 
 type Options struct {
 type Options struct {
@@ -104,6 +107,8 @@ func New(command []string, options *Options) (*App, error) {
 		},
 		},
 
 
 		titleTemplate: titleTemplate,
 		titleTemplate: titleTemplate,
+
+		onceMutex: umutex.New(),
 	}, nil
 	}, nil
 }
 }
 
 
@@ -314,6 +319,20 @@ func (app *App) handleWS(w http.ResponseWriter, r *http.Request) {
 			argv = append(argv, params...)
 			argv = append(argv, params...)
 		}
 		}
 	}
 	}
+
+	app.server.StartRoutine()
+
+	if app.options.Once {
+		if app.onceMutex.TryLock() { // no unlock required, it will die soon
+			log.Printf("Last client accepted, closing the listener.")
+			app.server.Close()
+		} else {
+			log.Printf("Server is already closing.")
+			conn.Close()
+			return
+		}
+	}
+
 	cmd := exec.Command(app.command[0], argv...)
 	cmd := exec.Command(app.command[0], argv...)
 	ptyIo, err := pty.Start(cmd)
 	ptyIo, err := pty.Start(cmd)
 	if err != nil {
 	if err != nil {

app/client_context.go

@@ -65,13 +65,6 @@ func (context *clientContext) goHandleClient() {
 		context.processReceive()
 		context.processReceive()
 	}()
 	}()
 
 
-	context.app.server.StartRoutine()
-
-	if context.app.options.Once {
-		log.Printf("Last client accepted, closing the listener.")
-		context.app.server.Close()
-	}
-
 	go func() {
 	go func() {
 		defer context.app.server.FinishRoutine()
 		defer context.app.server.FinishRoutine()