bash_go_service/backend-service/pkg/service/socket.go

package service

import (
	"bash_go_service/shared/pkg/client"
	"bash_go_service/shared/pkg/constants"
	"bash_go_service/shared/pkg/logger"
	"context"
	"encoding/binary"
	"net"
	"os"
	"path/filepath"
	"strings"
	"time"
)

type ExecParams struct {
	Filename string   `json:"filename"`
	Pwd      string   `json:"pwd"`
	Args     []string `json:"args"`
	Env      []string `json:"env"`
	LogPath  string   `json:"logPath"`
	Error    string   `json:"error"`
}

type UnixSocketTask struct {
	listener net.Listener
}

func NewUnixSocketTask() *UnixSocketTask {
	return &UnixSocketTask{}
}

func (t *UnixSocketTask) Execute(ctx context.Context) {
	socketPath := constants.SocketPath

	// 确保socket文件所在目录存在
	socketDir := filepath.Dir(socketPath)
	if _, err := os.Stat(socketDir); os.IsNotExist(err) {
		// 目录不存在，创建目录
		if err := os.MkdirAll(socketDir, 0755); err != nil {
			logger.Error("创建socket目录失败: %v\n", err)
			return
		}
	}

	// 删除已存在的socket文件
	os.Remove(socketPath)

	var err error
	t.listener, err = net.Listen("unix", socketPath)
	if err != nil {
		logger.Error("Listen error: %v\n", err)
		return
	}
	defer t.listener.Close()

	logger.Info("Unix socket service started, waiting for connections")

	// 创建用于优雅关闭的channel
	done := make(chan struct{})

	// 在新的goroutine中处理连接
	go func() {
		for {
			conn, err := t.listener.Accept()
			if err != nil {
				select {
				case <-ctx.Done():
					// 正常关闭
					close(done)
					return
				default:
					// 其他错误
					logger.Error("Accept error: %v\n", err)
					continue
				}
			}
			go t.handleConnection(conn)
		}
	}()

	// 等待上下文取消或服务完成
	select {
	case <-ctx.Done():
		// 关闭监听器触发Accept错误，从而退出接收循环
		t.listener.Close()
	case <-done:
		// 服务已经结束
	}

	logger.Info("Unix socket service stopped")
}

func (t *UnixSocketTask) Name() string {
	return "UnixSocketTask"
}

func (t *UnixSocketTask) writeMessage(conn net.Conn, message string) error {
	messageLen := uint32(len(message))
	if err := binary.Write(conn, binary.LittleEndian, messageLen); err != nil {
		return err
	}
	_, err := conn.Write([]byte(message))
	return err
}

func (t *UnixSocketTask) handleConnection(conn net.Conn) {
	defer conn.Close()

	params := &ExecParams{}

	// 读取文件名
	var filenameLen uint64
	binary.Read(conn, binary.LittleEndian, &filenameLen)
	filename := make([]byte, filenameLen)
	conn.Read(filename)
	params.Filename = string(filename)

	// 读取pwd
	var pwdLen uint64
	binary.Read(conn, binary.LittleEndian, &pwdLen)
	pwd := make([]byte, pwdLen)
	conn.Read(pwd)
	params.Pwd = string(pwd)

	// 读取argv
	var argc int32
	binary.Read(conn, binary.LittleEndian, &argc)
	params.Args = make([]string, argc)

	for i := 0; i < int(argc); i++ {
		var argLen uint64
		binary.Read(conn, binary.LittleEndian, &argLen)
		arg := make([]byte, argLen)
		conn.Read(arg)
		params.Args[i] = string(arg)
	}

	// 读取envp
	var envc int32
	binary.Read(conn, binary.LittleEndian, &envc)
	params.Env = make([]string, envc)

	for i := 0; i < int(envc); i++ {
		var envLen uint64
		binary.Read(conn, binary.LittleEndian, &envLen)
		env := make([]byte, envLen)
		conn.Read(env)
		params.Env[i] = string(env)
	}

	// 读取logPath
	var logPathLen uint64
	binary.Read(conn, binary.LittleEndian, &logPathLen)
	logPath := make([]byte, logPathLen)
	conn.Read(logPath)
	params.LogPath = string(logPath)

	context, readErr := os.ReadFile(params.LogPath)
	if readErr != nil {
		params.Error = ""
	} else {
		params.Error = string(context)
	}

	// 在HTTP调用前发送[start]标记
	// 在HTTP调用前发送[start]标记
	t.writeMessage(conn, "[sthttp]")

	initialMessage := constants.ColorBlue + "检测到您的命令执行出现错误! \n" + constants.ColorReset
	t.writeMessage(conn, initialMessage)
	logger.Debug("write: %s", initialMessage)
	// 显示加载动画
	loading := true
	go func() {
		frames := []string{"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"}
		i := 0
		for loading {
			t.writeMessage(conn, "\r正在查找解决方案"+frames[i%len(frames)])
			time.Sleep(100 * time.Millisecond)
			i++
		}

	}()

	time.Sleep(200 * time.Millisecond)
	client := client.NewClient()

	handleQuestion := func(question string) {
		if loading {
			loading = false // 停止加载动画
			t.writeMessage(conn, "\r")
		}

		// 替换控制标记
		question = strings.ReplaceAll(question, "[end]", "[/end]")
		question = strings.ReplaceAll(question, "[sthttp]", "[/sthttp]")

		// 当收到非空question时发送[res]响应
		res := constants.ColorGreen + question + constants.ColorBlue
		logger.Debug("write: %s", res)
		err := t.writeMessage(conn, res)
		if err != nil {
			logger.Debug("write error: %v", err)
			client.Destroy()
		}
	}

	err := client.PostToStream(constants.QuestionStreamApi, params, nil, handleQuestion)

	if err != nil {
		logger.Error("PostToStream error: %v", err)
	}

	// 结束后发送结束标记
	t.writeMessage(conn, "\n\r"+constants.ColorReset)
	t.writeMessage(conn, "[end]")
	conn.Close() // 关闭连接，防止客户端一直等待inf
	logger.Debug("connection done.")
}