//go:build !windows // +build !windows // Author: NiuJiuRu // Email: niujiuru@qq.com package shell import ( "bytes" "context" "fmt" "io" "os/exec" "syscall" "time" "github.com/mattn/go-shellwords" ) const ( gracePeriod = 2 * time.Second forceKillWait = 2 * time.Second exitTimeoutCode = 124 maxOutputSize = 1 << 20 // 最大 1 MB, 限制输出大小 checkProcessDelay = 50 * time.Millisecond ) type limitedBuffer struct { buf *bytes.Buffer limit int } func (l *limitedBuffer) Write(p []byte) (int, error) { remain := l.limit - l.buf.Len() if remain <= 0 { return len(p), nil } if len(p) > remain { p = p[:remain] } return l.buf.Write(p) } type processGroup struct { cmd *exec.Cmd pgid int } // 进程组是否存在 func (pg *processGroup) isProcessGroupAlive() bool { err := syscall.Kill(-pg.pgid, syscall.Signal(0)) if err == nil { return true } if errno, ok := err.(syscall.Errno); ok && errno == syscall.ESRCH { return false } return true } // 等待进程组终止 func (pg *processGroup) waitForTermination(timeout time.Duration) bool { deadline := time.Now().Add(timeout) for time.Now().Before(deadline) { if !pg.isProcessGroupAlive() { return true } time.Sleep(checkProcessDelay) } return !pg.isProcessGroupAlive() } // 终止整个进程组 func (pg *processGroup) terminate() error { if pg.cmd.Process == nil || pg.pgid <= 0 || !pg.isProcessGroupAlive() { return nil } // 第一阶段: 尝试优雅终止, SIGTERM if err := syscall.Kill(-pg.pgid, syscall.SIGTERM); err != nil { // 如果发送信号失败,可能进程已经不存在 if errno, ok := err.(syscall.Errno); ok && errno == syscall.ESRCH { return nil } } if pg.waitForTermination(gracePeriod) { return nil } // 第二阶段: 最后强制终止, SIGKILL if err := syscall.Kill(-pg.pgid, syscall.SIGKILL); err != nil { // 如果发送信号失败,可能进程已经不存在 if errno, ok := err.(syscall.Errno); ok && errno == syscall.ESRCH { return nil } return fmt.Errorf("failed to send SIGKILL to process group: %w", err) } if pg.waitForTermination(forceKillWait) { return nil } return fmt.Errorf("%w: process group %d still alive after force kill", ErrExecutorLostControl, pg.pgid) } func executeInternal(p ExecuteParams, onStart func(pg *processGroup)) (*ExecuteResult, error) { swp := shellwords.NewParser() swp.ParseEnv = true // 展开 "环境变量" swp.ParseBacktick = true // 展开 `...`命令 argv, err := swp.Parse(p.Cmd) if err != nil || len(argv) == 0 { return nil, ErrInvalidCommand } timeout := time.Duration(p.Timeout) * time.Second if timeout <= 0 { timeout = DefaultTimeout } ctx, cancel := context.WithTimeout(context.Background(), timeout) defer cancel() cmd := exec.Command(argv[0], argv[1:]...) if p.Dir != "" { // 设置工作目录 cmd.Dir = p.Dir } cmd.SysProcAttr = &syscall.SysProcAttr{ // 新的进程组 Setpgid: true, Pdeathsig: syscall.SIGKILL, } var stdout, stderr bytes.Buffer cmd.Stdout = io.Writer(&limitedBuffer{buf: &stdout, limit: maxOutputSize}) cmd.Stderr = io.Writer(&limitedBuffer{buf: &stderr, limit: maxOutputSize}) if err := cmd.Start(); err != nil { return nil, err } processInfo := &processGroup{ cmd: cmd, pgid: cmd.Process.Pid, // 进程组ID就是主进程的PID } if onStart != nil { onStart(processInfo) } done := make(chan error, 1) go func() { done <- cmd.Wait() }() exitCode := 0 var finalErr error select { case err := <-done: // 命令已结束, 输出结果 if err != nil { if ee, ok := err.(*exec.ExitError); ok { exitCode = ee.ExitCode() } else { return nil, err } } case <-ctx.Done(): /// 超时, kill整个进程组 exitCode = exitTimeoutCode if err := processInfo.terminate(); err != nil { finalErr = err break } timer := time.NewTimer(forceKillWait) defer timer.Stop() select { case <-done: case <-timer.C: finalErr = ErrExecutorLostControl } } return &ExecuteResult{ Stdout: stdout.String(), Stderr: stderr.String(), ExitCode: exitCode, }, finalErr } func Execute(p ExecuteParams) (*ExecuteResult, error) { return executeInternal(p, nil) }