From faf8300d8fc7be3be9d3f5e0eed31917a7a67116 Mon Sep 17 00:00:00 2001
From: Evan Chen <evanchen@kumoly.io>
Date: Wed, 24 Nov 2021 23:02:50 +0800
Subject: [PATCH] feat: #3

---
 gterm.go | 37 +++++++++++++++++++++++++++++++------
 1 file changed, 31 insertions(+), 6 deletions(-)

diff --git a/gterm.go b/gterm.go
index a016370..6f12675 100644
--- a/gterm.go
+++ b/gterm.go
@@ -30,10 +30,12 @@ var tmpl *engine.Engine
 var servePublic = http.FileServer(http.FS(public.FS))
 
 var pool map[string]chan struct{}
+var locks map[string]chan struct{}
 
 func init() {
 	tmpl = engine.Must(engine.New("").Parse(index))
 	pool = map[string]chan struct{}{}
+	locks = map[string]chan struct{}{}
 }
 
 type GTerm struct {
@@ -127,10 +129,11 @@ func (g *GTerm) App(w http.ResponseWriter, r *http.Request) {
 }
 
 type NewCmdRequest struct {
-	Cmd  string   `json:"cmd"`
-	Envs []string `json:"envs"`
-	Args []string `json:"args"`
-	Dir  string   `json:"dir"`
+	Cmd   string   `json:"cmd"`
+	Envs  []string `json:"envs"`
+	Args  []string `json:"args"`
+	Dir   string   `json:"dir"`
+	Block bool     `json:"block"`
 }
 
 func (g *GTerm) defaultCmd() *exec.Cmd {
@@ -142,6 +145,12 @@ func (g *GTerm) defaultCmd() *exec.Cmd {
 	return cmd
 }
 
+func (g *GTerm) echo(msg string) *exec.Cmd {
+	cmd := exec.Command("echo", msg)
+	cmd.Env = g.Envs
+	return cmd
+}
+
 func (g *GTerm) WS(w http.ResponseWriter, r *http.Request) {
 	id := fmt.Sprintf("[%02d] %s", ctr(), ksrv.GetIP(r))
 	pool[id] = make(chan struct{}, 1)
@@ -159,12 +168,28 @@ func (g *GTerm) WS(w http.ResponseWriter, r *http.Request) {
 		err := json.Unmarshal([]byte(param), req)
 		if err != nil {
 			l.Error().Err(err).Str("base", newCmd).Str("decrypt", param).Msg("cmd decode failed")
-			cmd = g.defaultCmd()
+			cmd = g.echo("cmd decode failed")
 		} else {
-			l.Info().Interface("cmd", req).Msg("starting cmd")
 			cmd = exec.Command(req.Cmd, req.Args...)
 			cmd.Env = append(req.Envs, "TERM=xterm-256color")
 			cmd.Dir = req.Dir
+			if req.Block {
+				lock, ok := locks[newCmd]
+				if !ok {
+					locks[newCmd] = make(chan struct{}, 1)
+				}
+				if len(lock) == 1 {
+					cmd = g.echo("cmd already running")
+				} else {
+					l.Info().Interface("cmd", req).Msg("starting cmd")
+					locks[newCmd] <- struct{}{}
+					defer func() {
+						<-locks[newCmd]
+						close(locks[newCmd])
+						delete(locks, newCmd)
+					}()
+				}
+			}
 		}
 	} else {
 		cmd = g.defaultCmd()