193 lines
3.7 KiB
Go
193 lines
3.7 KiB
Go
|
package main
|
||
|
|
||
|
import (
|
||
|
"encoding/json"
|
||
|
"fmt"
|
||
|
"net"
|
||
|
"os"
|
||
|
"strings"
|
||
|
|
||
|
"github.com/google/uuid"
|
||
|
)
|
||
|
|
||
|
type CmdDaemon struct {
|
||
|
SocketPath string // unix socket path
|
||
|
cmds map[string]CmdHandler
|
||
|
}
|
||
|
type CmdRequest struct {
|
||
|
Id string `json:"id"`
|
||
|
Cmd string `json:"cmd"`
|
||
|
Args string `json:"args"`
|
||
|
IsDebug bool `json:"debug"`
|
||
|
}
|
||
|
type CmdResponse struct {
|
||
|
Id string `json:"id"`
|
||
|
Data string `json:"data"`
|
||
|
Error string `json:"error"`
|
||
|
Continue bool `json:"continue"`
|
||
|
}
|
||
|
|
||
|
type CmdConn struct {
|
||
|
net.Conn
|
||
|
Id string
|
||
|
}
|
||
|
|
||
|
func (c *CmdConn) Write(d string) error {
|
||
|
resp := CmdResponse{
|
||
|
Id: c.Id,
|
||
|
Data: d,
|
||
|
Continue: true,
|
||
|
}
|
||
|
return Write(c.Conn, resp)
|
||
|
}
|
||
|
func (c *CmdConn) WriteError(err error, isContinue bool) error {
|
||
|
resp := CmdResponse{
|
||
|
Id: c.Id,
|
||
|
Error: err.Error(),
|
||
|
Continue: isContinue,
|
||
|
}
|
||
|
return Write(c.Conn, resp)
|
||
|
}
|
||
|
func (c *CmdConn) End(d string) error {
|
||
|
resp := CmdResponse{
|
||
|
Id: c.Id,
|
||
|
Data: d,
|
||
|
Continue: false,
|
||
|
}
|
||
|
return Write(c.Conn, resp)
|
||
|
}
|
||
|
|
||
|
type CmdHandler interface {
|
||
|
Handle(conn *CmdConn, req *CmdRequest) error
|
||
|
Description() string
|
||
|
Usage() string
|
||
|
}
|
||
|
|
||
|
func (c *CmdDaemon) Listen() error {
|
||
|
// 删除已存在的 socket 文件(如果存在)
|
||
|
if err := os.Remove(c.SocketPath); err != nil && !os.IsNotExist(err) {
|
||
|
return fmt.Errorf("failed to remove existing socket file: %v", err)
|
||
|
}
|
||
|
|
||
|
// 监听 unix socket
|
||
|
listener, err := net.Listen("unix", c.SocketPath)
|
||
|
if err != nil {
|
||
|
return fmt.Errorf("failed to listen on socket: %v", err)
|
||
|
}
|
||
|
defer listener.Close()
|
||
|
|
||
|
// 设置 socket 文件权限
|
||
|
if err := os.Chmod(c.SocketPath, 0777); err != nil {
|
||
|
return fmt.Errorf("failed to set socket file permissions: %v", err)
|
||
|
}
|
||
|
|
||
|
for {
|
||
|
conn, err := listener.Accept()
|
||
|
if err != nil {
|
||
|
return fmt.Errorf("failed to accept connection: %v", err)
|
||
|
}
|
||
|
|
||
|
// 处理每个连接
|
||
|
go func(conn net.Conn) {
|
||
|
defer conn.Close()
|
||
|
|
||
|
req, err := Read[CmdRequest](conn)
|
||
|
if err != nil {
|
||
|
_ = Write(conn, CmdResponse{
|
||
|
Error: "failed to read request: " + err.Error(),
|
||
|
Continue: false,
|
||
|
})
|
||
|
return
|
||
|
}
|
||
|
|
||
|
cmdHandler, ok := c.cmds[req.Cmd]
|
||
|
if !ok {
|
||
|
_ = Write(conn, CmdResponse{
|
||
|
Error: "unknown command: " + req.Cmd,
|
||
|
Continue: false,
|
||
|
})
|
||
|
return
|
||
|
}
|
||
|
|
||
|
// 执行命令处理程序
|
||
|
cmdConn := &CmdConn{Conn: conn, Id: req.Id}
|
||
|
err = cmdHandler.Handle(cmdConn, req)
|
||
|
if err != nil {
|
||
|
_ = cmdConn.WriteError(err, false)
|
||
|
}
|
||
|
}(conn)
|
||
|
}
|
||
|
}
|
||
|
|
||
|
func (c *CmdDaemon) RegisterCmd(cmd string, handler CmdHandler) {
|
||
|
c.cmds[cmd] = handler
|
||
|
}
|
||
|
|
||
|
func (c *CmdDaemon) isDebug(cmd string) bool {
|
||
|
|
||
|
return cmd == "--debug" || cmd == "-d"
|
||
|
}
|
||
|
func (c *CmdDaemon) Run() error {
|
||
|
// 从命令参数中解析出是否debug 子命令和剩余参数字符串
|
||
|
args := os.Args[1:]
|
||
|
isDebug := c.isDebug(args[0])
|
||
|
var remainingArgs []string
|
||
|
cmd := ""
|
||
|
if isDebug {
|
||
|
if len(args) > 1 {
|
||
|
cmd = args[1]
|
||
|
remainingArgs = args[2:]
|
||
|
} else {
|
||
|
cmd = "help"
|
||
|
isDebug = false
|
||
|
remainingArgs = []string{}
|
||
|
}
|
||
|
} else {
|
||
|
if len(args) > 0 {
|
||
|
cmd = args[0]
|
||
|
remainingArgs = args[1:]
|
||
|
} else {
|
||
|
cmd = "help"
|
||
|
isDebug = false
|
||
|
remainingArgs = []string{}
|
||
|
}
|
||
|
}
|
||
|
cmdReq := CmdRequest{
|
||
|
Args: strings.Join(remainingArgs, " "),
|
||
|
Cmd: cmd,
|
||
|
Id: uuid.New().String(),
|
||
|
IsDebug: isDebug,
|
||
|
}
|
||
|
// dial unix socket
|
||
|
conn, err := net.Dial("unix", c.SocketPath)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
defer conn.Close()
|
||
|
// send cmd request
|
||
|
cmdReqJson, err := json.Marshal(cmdReq)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
_, err = conn.Write(cmdReqJson)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
|
||
|
for {
|
||
|
|
||
|
resp, err := Read[CmdResponse](conn)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
if resp.Error != "" {
|
||
|
fmt.Println(resp.Error)
|
||
|
}
|
||
|
if !resp.Continue {
|
||
|
break
|
||
|
}
|
||
|
}
|
||
|
return nil
|
||
|
|
||
|
}
|