From eddd813846900bac13f7a3aeef6c99d27ea81d1e Mon Sep 17 00:00:00 2001
From: Evan Chen <evanchen@kumoly.io>
Date: Thu, 11 Nov 2021 09:58:36 +0800
Subject: [PATCH] limit run command to single instance

---
 configui.go |  7 +++++--
 file.go     | 36 ++++++++++++++++++++++++++++--------
 handler.go  | 33 +++++++++++++++++++++++++++++++--
 3 files changed, 64 insertions(+), 12 deletions(-)

diff --git a/configui.go b/configui.go
index e0cd085..54d2ed7 100644
--- a/configui.go
+++ b/configui.go
@@ -34,14 +34,17 @@ var Ext2Mode map[string]string = map[string]string{
 }
 
 type Action struct {
-	Name string `json:"name"`
-	Cmd  string `json:"cmd"`
+	Name string        `json:"name"`
+	Cmd  string        `json:"cmd"`
+	run  chan struct{} `json:"-"`
+	pid  int           `json:"-"`
 }
 
 type Integration struct {
 	Name        string                 `json:"name"`
 	Description string                 `json:"description"`
 	Cmd         func() (string, error) `json:"-"`
+	run         chan struct{}          `json:"-"`
 }
 
 type ConfigUI struct {
diff --git a/file.go b/file.go
index 38616c2..b4e21ea 100644
--- a/file.go
+++ b/file.go
@@ -1,7 +1,9 @@
 package configui
 
 import (
+	"bytes"
 	"errors"
+	"fmt"
 	"os"
 	"os/exec"
 	"runtime"
@@ -23,8 +25,10 @@ type File struct {
 	// used for parsing post data
 	Data string `json:"data,omitempty"`
 
-	lock  sync.RWMutex `json:"-"`
-	owner *ConfigUI    `json:"-"`
+	lock  sync.RWMutex  `json:"-"`
+	owner *ConfigUI     `json:"-"`
+	run   chan struct{} `json:"-"`
+	pid   int           `json:"-"`
 }
 
 func (f *File) Read() ([]byte, error) {
@@ -57,6 +61,20 @@ func (f *File) Do(CmdTimeout time.Duration) (string, error) {
 	}
 	f.lock.RLock()
 	defer f.lock.RUnlock()
+
+	// limit running instance to 1
+	if cap(f.run) == 0 {
+		f.run = make(chan struct{}, 1)
+	}
+	select {
+	case f.run <- struct{}{}:
+		defer func() { <-f.run }()
+	default:
+		f.owner.log.Error("task is running: ", f.Name)
+		return "", fmt.Errorf("another task of %s is running with pid: %d", f.Name, f.pid)
+	}
+
+	// prepare cmd
 	cmd := &exec.Cmd{}
 	if runtime.GOOS == "windows" {
 		cmd = exec.Command(WIN_SHELL, "/c", f.Action)
@@ -66,12 +84,14 @@ func (f *File) Do(CmdTimeout time.Duration) (string, error) {
 	f.owner.log.Info("DO: ", f.Action)
 	done := make(chan string, 1)
 	go func() {
-		out, _ := cmd.CombinedOutput()
-		// real cmd err is unhandled, but passed to client
-		// if err != nil {
-		// 	return string(out), err
-		// }
-		done <- string(out)
+
+		var b bytes.Buffer
+		cmd.Stdout = &b
+		cmd.Stderr = &b
+		cmd.Start()
+		f.pid = cmd.Process.Pid
+		cmd.Wait()
+		done <- b.String()
 	}()
 	select {
 	case <-time.After(CmdTimeout):
diff --git a/handler.go b/handler.go
index c593c45..e81c71e 100644
--- a/handler.go
+++ b/handler.go
@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 	"strconv"
+	"time"
 
 	"kumoly.io/lib/ksrv"
 )
@@ -130,9 +131,25 @@ func (cui *ConfigUI) Apply(w http.ResponseWriter, r *http.Request) {
 
 func (cui *ConfigUI) DoAction(w http.ResponseWriter, r *http.Request) {
 	name := r.URL.Query().Get("name")
-	for _, v := range cui.Actions {
+	for i, v := range cui.Actions {
 		if v.Name == name {
+
+			// limit running instance to one
+			if cap(cui.Actions[i].run) != 1 {
+				cui.Actions[i].run = make(chan struct{}, 1)
+			}
+			select {
+			case cui.Actions[i].run <- struct{}{}:
+				defer func() { <-cui.Actions[i].run }()
+			default:
+				panic(fmt.Errorf("another task of %s is running with pid: %d", name, v.pid))
+			}
+
 			file := &File{Name: name, Action: v.Cmd, owner: cui}
+			go func() {
+				<-time.After(time.Millisecond * 10)
+				cui.Actions[i].pid = file.pid
+			}()
 			result, err := file.Do(cui.cmdTimeout)
 			if err != nil {
 				panic(err)
@@ -146,8 +163,20 @@ func (cui *ConfigUI) DoAction(w http.ResponseWriter, r *http.Request) {
 
 func (cui *ConfigUI) DoIntegration(w http.ResponseWriter, r *http.Request) {
 	name := r.URL.Query().Get("name")
-	for _, v := range cui.Integrations {
+	for i, v := range cui.Integrations {
 		if v.Name == name {
+
+			// limit running instance to one
+			if cap(cui.Integrations[i].run) != 1 {
+				cui.Integrations[i].run = make(chan struct{}, 1)
+			}
+			select {
+			case cui.Integrations[i].run <- struct{}{}:
+				defer func() { <-cui.Integrations[i].run }()
+			default:
+				panic(fmt.Errorf("another task of %s is running", name))
+			}
+
 			result, err := v.Cmd()
 			if err != nil {
 				panic(err)