1Panel/core/utils/cmd/cmdx.go

package cmd

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/1Panel-dev/1Panel/core/app/task"
	"github.com/1Panel-dev/1Panel/core/buserr"
	"github.com/1Panel-dev/1Panel/core/constant"
)

const maxStreamOutputCapture = 64 * 1024

type CommandHelper struct {
	context      context.Context
	workDir      string
	outputFile   string
	env          []string
	timeout      time.Duration
	taskItem     *task.Task
	logger       *log.Logger
	IgnoreExist1 bool
}

type Option func(*CommandHelper)

type PipeCommand struct {
	Name  string
	Args  []string
	Env   []string
	Dir   string
	Stdin io.Reader
}

type lockedBuffer struct {
	mu        sync.Mutex
	buf       bytes.Buffer
	limit     int
	truncated int
}

func (b *lockedBuffer) Write(p []byte) (int, error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.limit > 0 && b.buf.Len() >= b.limit {
		b.truncated += len(p)
		return len(p), nil
	}
	if b.limit > 0 && b.buf.Len()+len(p) > b.limit {
		keep := b.limit - b.buf.Len()
		_, _ = b.buf.Write(p[:keep])
		b.truncated += len(p) - keep
		return len(p), nil
	}
	return b.buf.Write(p)
}

func (b *lockedBuffer) String() string {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.truncated == 0 {
		return b.buf.String()
	}
	return fmt.Sprintf("%s\n... truncated %d bytes ...", b.buf.String(), b.truncated)
}

func NewCommandMgr(opts ...Option) *CommandHelper {
	s := &CommandHelper{}
	for _, opt := range opts {
		opt(s)
	}
	return s
}

func (c *CommandHelper) Run(name string, arg ...string) error {
	_, err := c.run(name, arg...)
	return err
}

func (c *CommandHelper) RunWithStdout(name string, arg ...string) (string, error) {
	return c.run(name, arg...)
}

func (c *CommandHelper) RunPipe(commands ...PipeCommand) (string, error) {
	if len(commands) == 0 {
		return "", nil
	}

	ctx, cancel, cmds := c.preparePipeCommands(commands)
	if cancel != nil {
		defer cancel()
	}

	customWriter := &CustomWriter{taskItem: c.taskItem}
	var outputFile *os.File
	stdout, stderr := &lockedBuffer{}, &lockedBuffer{}
	limitOutputCapture := c.taskItem != nil || c.logger != nil || len(c.outputFile) != 0
	if limitOutputCapture {
		stdout.limit = maxStreamOutputCapture
		stderr.limit = maxStreamOutputCapture
	}
	var pipeStderr io.Writer = stderr
	var lastStdout io.Writer = stdout
	var lastStderr io.Writer = stderr
	var streamWriter io.Writer
	if c.taskItem != nil {
		streamWriter = customWriter
	} else if c.logger != nil {
		streamWriter = c.logger.Writer()
	} else if len(c.outputFile) != 0 {
		file, err := os.OpenFile(c.outputFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, constant.FilePerm)
		if err != nil {
			return "", err
		}
		outputFile = file
		lastStdout = outputFile
	}
	if streamWriter != nil {
		pipeStderr = io.MultiWriter(stderr, streamWriter)
		lastStdout = io.MultiWriter(stdout, streamWriter)
		lastStderr = io.MultiWriter(stderr, streamWriter)
	}
	defer func() {
		if c.taskItem != nil {
			customWriter.Flush()
		}
		if closer, ok := streamWriter.(io.Closer); ok {
			_ = closer.Close()
		}
		if outputFile != nil {
			_ = outputFile.Close()
		}
	}()
	if err := connectPipeCommands(cmds, lastStdout, lastStderr, pipeStderr); err != nil {
		return "", err
	}
	if err := startPipeCommands(cmds); err != nil {
		return handleErrString(stdout.String(), stderr.String(), c.IgnoreExist1, err)
	}

	runErr := c.pipeResultErr(ctx, waitPipeCommands(ctx, cmds))
	if runErr != nil {
		return handleErrString(stdout.String(), stderr.String(), c.IgnoreExist1, runErr)
	}
	return stdout.String(), nil
}

func (c *CommandHelper) RunPipeToFile(outputFile string, commands ...PipeCommand) (string, error) {
	if len(commands) == 0 {
		return "", nil
	}

	ctx, cancel, cmds := c.preparePipeCommands(commands)
	if cancel != nil {
		defer cancel()
	}

	file, err := os.OpenFile(outputFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, constant.FilePerm)
	if err != nil {
		return "", err
	}
	defer func() { _ = file.Close() }()

	stderr := &lockedBuffer{limit: maxStreamOutputCapture}
	if err := connectPipeCommands(cmds, file, stderr, stderr); err != nil {
		return "", err
	}
	if err := startPipeCommands(cmds); err != nil {
		return handleErrString("", stderr.String(), c.IgnoreExist1, err)
	}

	runErr := c.pipeResultErr(ctx, waitPipeCommands(ctx, cmds))
	if runErr != nil {
		return handleErrString("", stderr.String(), c.IgnoreExist1, runErr)
	}
	return "", nil
}

func (c *CommandHelper) preparePipeCommands(commands []PipeCommand) (context.Context, context.CancelFunc, []*exec.Cmd) {
	ctx, cancel := c.pipeContext()
	cmds := c.buildPipeCommands(ctx, commands)
	if commands[0].Stdin != nil {
		cmds[0].Stdin = commands[0].Stdin
	}
	return ctx, cancel, cmds
}

func (c *CommandHelper) pipeResultErr(ctx context.Context, runErr error) error {
	if errors.Is(ctx.Err(), context.DeadlineExceeded) {
		return buserr.New("ErrCmdTimeout")
	}
	if errors.Is(ctx.Err(), context.Canceled) {
		return buserr.New("ErrShutDown")
	}
	return runErr
}

func (c *CommandHelper) pipeContext() (context.Context, context.CancelFunc) {
	ctx := c.context
	if ctx == nil {
		ctx = context.Background()
	}
	if c.timeout == 0 {
		return ctx, nil
	}
	return context.WithTimeout(ctx, c.timeout)
}

func (c *CommandHelper) buildPipeCommands(ctx context.Context, commands []PipeCommand) []*exec.Cmd {
	cmds := make([]*exec.Cmd, 0, len(commands))
	for _, item := range commands {
		cmdItem := exec.CommandContext(ctx, item.Name, item.Args...)
		cmdItem.Env = append(os.Environ(), c.env...)
		cmdItem.Env = append(cmdItem.Env, item.Env...)
		cmdItem.Dir = c.workDir
		if item.Dir != "" {
			cmdItem.Dir = item.Dir
		}
		cmdItem.SysProcAttr = &syscall.SysProcAttr{
			Setpgid: true,
		}
		cmds = append(cmds, cmdItem)
	}
	return cmds
}

func connectPipeCommands(cmds []*exec.Cmd, stdout, stderr, pipeStderr io.Writer) error {
	for i := 0; i < len(cmds)-1; i++ {
		pipe, err := cmds[i].StdoutPipe()
		if err != nil {
			return err
		}
		cmds[i+1].Stdin = pipe
		cmds[i].Stderr = pipeStderr
	}
	last := cmds[len(cmds)-1]
	last.Stdout = stdout
	last.Stderr = stderr
	return nil
}

func startPipeCommands(cmds []*exec.Cmd) error {
	for i := len(cmds) - 1; i >= 0; i-- {
		if err := cmds[i].Start(); err != nil {
			killStarted(cmds[i+1:])
			return err
		}
	}
	return nil
}

func waitPipeCommands(ctx context.Context, cmds []*exec.Cmd) error {
	done := make(chan error, 1)
	go func() {
		var runErr error
		for _, item := range cmds {
			if err := item.Wait(); err != nil && runErr == nil {
				runErr = err
			}
		}
		done <- runErr
	}()
	select {
	case runErr := <-done:
		return runErr
	case <-ctx.Done():
		killProcessGroups(cmds)
		return <-done
	}
}

func (c *CommandHelper) run(name string, arg ...string) (string, error) {
	var cmd *exec.Cmd
	var newContext context.Context
	var cancel context.CancelFunc
	var outputFile *os.File

	if c.timeout != 0 {
		if c.context == nil {
			newContext, cancel = context.WithTimeout(context.Background(), c.timeout)
		} else {
			newContext, cancel = context.WithTimeout(c.context, c.timeout)
		}
		defer cancel()
	} else if c.context != nil {
		newContext = c.context
	}

	if newContext != nil {
		cmd = exec.CommandContext(newContext, name, arg...)
	} else {
		cmd = exec.Command(name, arg...)
	}
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setpgid: true,
	}

	customWriter := &CustomWriter{taskItem: c.taskItem}
	var stdout, stderr bytes.Buffer
	var loggerCloser io.Closer
	if c.taskItem != nil {
		cmd.Stdout = io.MultiWriter(&stdout, customWriter)
		cmd.Stderr = io.MultiWriter(&stderr, customWriter)
	} else if c.logger != nil {
		streamWriter := c.logger.Writer()
		if closer, ok := streamWriter.(io.Closer); ok {
			loggerCloser = closer
		}
		cmd.Stdout = io.MultiWriter(&stdout, streamWriter)
		cmd.Stderr = io.MultiWriter(&stderr, streamWriter)
	} else if len(c.outputFile) != 0 {
		file, err := os.OpenFile(c.outputFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, constant.FilePerm)
		if err != nil {
			return "", err
		}
		outputFile = file
		cmd.Stdout = io.MultiWriter(&stdout, outputFile)
		cmd.Stderr = io.MultiWriter(&stderr, outputFile)
	} else {
		cmd.Stdout = &stdout
		cmd.Stderr = &stderr
	}
	cmd.Env = append(os.Environ(), c.env...)
	if len(c.workDir) != 0 {
		cmd.Dir = c.workDir
	}
	defer func() {
		if loggerCloser != nil {
			_ = loggerCloser.Close()
		}
		if outputFile != nil {
			_ = outputFile.Close()
		}
	}()

	if err := cmd.Start(); err != nil {
		return "", fmt.Errorf("cmd start failed: %w", err)
	}
	if c.taskItem != nil {
		defer customWriter.Flush()
	}

	done := make(chan error, 1)
	go func() {
		done <- cmd.Wait()
	}()
	select {
	case err := <-done:
		if err != nil {
			return handleErr(&stdout, &stderr, c.IgnoreExist1, err)
		}
		return stdout.String(), nil
	case <-contextDone(newContext):
		if cmd.Process != nil && cmd.Process.Pid > 0 {
			_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
		}
		var err error
		switch newContext.Err() {
		case context.DeadlineExceeded:
			err = buserr.New("ErrCmdTimeout")
		case context.Canceled:
			err = buserr.New("ErrShutDown")
		default:
			err = newContext.Err()
		}
		<-done
		return "", err
	}
}

func contextDone(ctx context.Context) <-chan struct{} {
	if ctx == nil {
		return nil
	}
	return ctx.Done()
}

func killStarted(cmds []*exec.Cmd) {
	killProcessGroups(cmds)
	for _, item := range cmds {
		if item.Process != nil {
			_ = item.Wait()
		}
	}
}

func killProcessGroups(cmds []*exec.Cmd) {
	for _, item := range cmds {
		if item.Process != nil {
			_ = syscall.Kill(-item.Process.Pid, syscall.SIGKILL)
		}
	}
}

func WithOutputFile(outputFile string) Option {
	return func(s *CommandHelper) {
		s.outputFile = outputFile
	}
}

func WithContext(ctx context.Context) Option {
	return func(s *CommandHelper) {
		s.context = ctx
	}
}

func WithTimeout(timeout time.Duration) Option {
	return func(s *CommandHelper) {
		s.timeout = timeout
	}
}

func WithLogger(logger *log.Logger) Option {
	return func(s *CommandHelper) {
		s.logger = logger
	}
}

func WithTask(taskItem task.Task) Option {
	return func(s *CommandHelper) {
		s.taskItem = &taskItem
	}
}

func WithWorkDir(workDir string) Option {
	return func(s *CommandHelper) {
		s.workDir = workDir
	}
}

func WithEnv(env ...string) Option {
	return func(s *CommandHelper) {
		s.env = append(s.env, env...)
	}
}

func WithIgnoreExist1() Option {
	return func(s *CommandHelper) {
		s.IgnoreExist1 = true
	}
}

type CustomWriter struct {
	mu       sync.Mutex
	taskItem *task.Task
	buffer   bytes.Buffer
}

func (cw *CustomWriter) Write(p []byte) (n int, err error) {
	cw.mu.Lock()
	defer cw.mu.Unlock()
	cw.buffer.Write(p)
	lines := strings.Split(cw.buffer.String(), "\n")

	for i := 0; i < len(lines)-1; i++ {
		cw.taskItem.Log(lines[i])
	}
	cw.buffer.Reset()
	cw.buffer.WriteString(lines[len(lines)-1])

	return len(p), nil
}

func (cw *CustomWriter) Flush() {
	cw.mu.Lock()
	defer cw.mu.Unlock()
	if cw.buffer.Len() > 0 {
		cw.taskItem.Log(cw.buffer.String())
		cw.buffer.Reset()
	}
}

func handleErr(stdout, stderr fmt.Stringer, ignoreExist1 bool, err error) (string, error) {
	return handleErrString(stdout.String(), stderr.String(), ignoreExist1, err)
}

func handleErrString(stdout, stderr string, ignoreExist1 bool, err error) (string, error) {
	var exitError *exec.ExitError
	if ignoreExist1 && errors.As(err, &exitError) {
		if status, ok := exitError.Sys().(syscall.WaitStatus); ok {
			if status.ExitStatus() == 1 {
				return "", nil
			}
		}
	}
	outItem := stdout
	errItem := stderr
	if len(errItem) != 0 && len(outItem) != 0 {
		return outItem, fmt.Errorf("stdout: %s; stderr: %s, err: %v", outItem, errItem, err)
	}
	if len(errItem) != 0 {
		return outItem, fmt.Errorf("stderr: %s, err: %v", errItem, err)
	}
	if len(outItem) != 0 {
		return outItem, fmt.Errorf("stdout: %s, err: %v", outItem, err)
	}
	return "", err
}