From 7422de126939a1242554eda57fd287b05c2b3390 Mon Sep 17 00:00:00 2001
From: naiba <hi@nai.ba>
Date: Thu, 30 Sep 2021 12:01:01 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20=E4=BF=AE=E5=A4=8D=20WebSSH=20?=
 =?UTF-8?q?=E8=BF=9E=E6=8E=A5=E6=96=AD=E5=BC=80=E5=90=8E=E9=81=97=E7=95=99?=
 =?UTF-8?q?=E5=83=B5=E5=B0=B8=E8=BF=9B=E7=A8=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/agent/main.go                       |  5 ++--
 cmd/agent/monitor/monitor.go            |  4 ++-
 cmd/agent/processgroup/process_group.go | 35 ++++++++++++++++++++-----
 cmd/agent/pty/pty.go                    | 16 +++++++++--
 4 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/cmd/agent/main.go b/cmd/agent/main.go
index 896531e..0df92e1 100644
--- a/cmd/agent/main.go
+++ b/cmd/agent/main.go
@@ -330,6 +330,7 @@ func handleCommandTask(task *pb.Task, result *pb.TaskResult) {
 		result.Data = string(output)
 		result.Successful = true
 	}
+	pg.Dispose()
 	result.Delay = float32(time.Since(startedAt).Seconds())
 }
 
@@ -365,9 +366,9 @@ func handleTerminalTask(task *pb.Task) {
 	}
 
 	defer func() {
-		tty.Close()
+		err := tty.Close()
 		conn.Close()
-		println("terminal exit", terminal.Session)
+		println("terminal exit", terminal.Session, err)
 	}()
 	println("terminal init", terminal.Session)
 
diff --git a/cmd/agent/monitor/monitor.go b/cmd/agent/monitor/monitor.go
index b330f4a..5f79d28 100644
--- a/cmd/agent/monitor/monitor.go
+++ b/cmd/agent/monitor/monitor.go
@@ -64,7 +64,9 @@ func GetHost() *model.Host {
 		swapMemTotal = mv.SwapTotal
 	}
 
-	cachedBootTime = time.Unix(int64(hi.BootTime), 0)
+	if cachedBootTime.IsZero() {
+		cachedBootTime = time.Unix(int64(hi.BootTime), 0)
+	}
 
 	return &model.Host{
 		Platform:        hi.OS,
diff --git a/cmd/agent/processgroup/process_group.go b/cmd/agent/processgroup/process_group.go
index 83195c9..2bda9bd 100644
--- a/cmd/agent/processgroup/process_group.go
+++ b/cmd/agent/processgroup/process_group.go
@@ -1,9 +1,11 @@
+//go:build !windows
 // +build !windows
 
 package processgroup
 
 import (
 	"os/exec"
+	"sync"
 	"syscall"
 )
 
@@ -15,13 +17,34 @@ func NewProcessExitGroup() (ProcessExitGroup, error) {
 	return ProcessExitGroup{}, nil
 }
 
-func (g *ProcessExitGroup) Dispose() error {
-	for _, c := range g.cmds {
-		if err := syscall.Kill(-c.Process.Pid, syscall.SIGKILL); err != nil {
-			return err
-		}
+func (g *ProcessExitGroup) killChildProcess(c *exec.Cmd) error {
+	pgid, err := syscall.Getpgid(c.Process.Pid)
+	if err != nil {
+		// Fall-back on error. Kill the main process only.
+		c.Process.Kill()
 	}
-	return nil
+	// Kill the whole process group.
+	syscall.Kill(-pgid, syscall.SIGTERM)
+	return c.Wait()
+}
+
+func (g *ProcessExitGroup) Dispose() []error {
+	var errors []error
+	mutex := new(sync.Mutex)
+	wg := new(sync.WaitGroup)
+	wg.Add(len(g.cmds))
+	for _, c := range g.cmds {
+		go func(c *exec.Cmd) {
+			defer wg.Done()
+			if err := g.killChildProcess(c); err != nil {
+				mutex.Lock()
+				defer mutex.Unlock()
+				errors = append(errors, err)
+			}
+		}(c)
+	}
+	wg.Wait()
+	return errors
 }
 
 func (g *ProcessExitGroup) AddProcess(cmd *exec.Cmd) error {
diff --git a/cmd/agent/pty/pty.go b/cmd/agent/pty/pty.go
index 2e8fc58..c9841fc 100644
--- a/cmd/agent/pty/pty.go
+++ b/cmd/agent/pty/pty.go
@@ -1,5 +1,5 @@
 //go:build !windows
-//+build !windows
+// +build !windows
 
 package pty
 
@@ -7,6 +7,7 @@ import (
 	"errors"
 	"os"
 	"os/exec"
+	"syscall"
 
 	opty "github.com/creack/pty"
 )
@@ -53,9 +54,20 @@ func (pty *Pty) Setsize(cols, rows uint32) error {
 	})
 }
 
+func (pty *Pty) killChildProcess(c *exec.Cmd) error {
+	pgid, err := syscall.Getpgid(c.Process.Pid)
+	if err != nil {
+		// Fall-back on error. Kill the main process only.
+		c.Process.Kill()
+	}
+	// Kill the whole process group.
+	syscall.Kill(-pgid, syscall.SIGTERM)
+	return c.Wait()
+}
+
 func (pty *Pty) Close() error {
 	if err := pty.tty.Close(); err != nil {
 		return err
 	}
-	return pty.cmd.Process.Kill()
+	return pty.killChildProcess(pty.cmd)
 }