modularize monitor, reduce init usage (#81)

* kill process tree using syscall on windows & cleanup (#80)

* kill process tree using syscall on windows & cleanup

* use job api

* add error check for cmd.Start

* modularize monitor, reduce init usage

* replace slices with sort

* update gopsutil & other dependencies
This commit is contained in:
UUBulb
2024-11-03 21:53:09 +08:00
committed by GitHub
parent 0cba96bae1
commit af41e4d843
23 changed files with 665 additions and 381 deletions
+76 -67
View File
@@ -1,6 +1,7 @@
package main
import (
"bytes"
"context"
"crypto/md5"
"crypto/tls"
@@ -12,7 +13,6 @@ import (
"net/http"
"net/url"
"os"
"os/exec"
"path/filepath"
"runtime"
"strings"
@@ -45,10 +45,10 @@ import (
)
var (
version string
version = monitor.Version // 来自于 GoReleaser 的版本号
arch string
defaultConfigPath string
executablePath string
defaultConfigPath = loadDefaultConfigPath()
client pb.NezhaServiceClient
initialized bool
dnsResolver = &net.Resolver{PreferGo: true}
@@ -73,7 +73,7 @@ const (
networkTimeOut = time.Second * 5 // 普通网络超时
)
func init() {
func setEnv() {
resolver.SetDefaultScheme("passthrough")
net.DefaultResolver.PreferGo = true // 使用 Go 内置的 DNS 解析器解析域名
net.DefaultResolver.Dial = func(ctx context.Context, network, address string) (net.Conn, error) {
@@ -102,17 +102,67 @@ func init() {
utls.HelloChrome_Auto, new(utls.Config),
http.DefaultTransport, nil, &headers,
)
}
// 来自于 GoReleaser 的版本号
monitor.Version = version
func loadDefaultConfigPath() string {
var err error
executablePath, err = os.Executable()
if err != nil {
panic(err)
return ""
}
return filepath.Join(filepath.Dir(executablePath), "config.yml")
}
func preRun(configPath string) error {
// init
setEnv()
if configPath == "" {
configPath = defaultConfigPath
}
defaultConfigPath = filepath.Join(filepath.Dir(executablePath), "config.yml")
// windows环境处理
if runtime.GOOS == "windows" {
hostArch, err := host.KernelArch()
if err != nil {
return err
}
switch hostArch {
case "i386", "i686":
hostArch = "386"
case "x86_64":
hostArch = "amd64"
case "aarch64":
hostArch = "arm64"
}
if arch != hostArch {
return fmt.Errorf("与当前系统不匹配,当前运行 %s_%s, 需要下载 %s_%s", runtime.GOOS, arch, runtime.GOOS, hostArch)
}
}
if err := agentConfig.Read(configPath); err != nil {
return fmt.Errorf("打开配置文件失败:%v", err)
}
monitor.InitConfig(&agentConfig)
if agentConfig.ClientSecret == "" {
return errors.New("ClientSecret 不能为空")
}
if agentConfig.ReportDelay < 1 || agentConfig.ReportDelay > 4 {
return errors.New("report-delay 的区间为 1-4")
}
if agentConfig.UUID == "" {
if uuid, err := uuid.GenerateUUID(); err == nil {
agentConfig.UUID = uuid
return agentConfig.Save()
} else {
return fmt.Errorf("生成 UUID 失败:%v", err)
}
}
return nil
}
func main() {
@@ -129,9 +179,13 @@ func main() {
return nil
}
if path := c.String("config"); path != "" {
preRun(path)
if err := preRun(path); err != nil {
return err
}
} else {
preRun(defaultConfigPath)
if err := preRun(""); err != nil {
return err
}
}
runService("", "")
return nil
@@ -181,51 +235,6 @@ func main() {
}
}
func preRun(configPath string) {
// windows环境处理
if runtime.GOOS == "windows" {
hostArch, err := host.KernelArch()
if err != nil {
panic(err)
}
if hostArch == "i386" {
hostArch = "386"
}
if hostArch == "i686" || hostArch == "ia64" || hostArch == "x86_64" {
hostArch = "amd64"
}
if hostArch == "aarch64" {
hostArch = "arm64"
}
if arch != hostArch {
panic(fmt.Sprintf("与当前系统不匹配,当前运行 %s_%s, 需要下载 %s_%s", runtime.GOOS, arch, runtime.GOOS, hostArch))
}
}
if err := agentConfig.Read(configPath); err != nil {
log.Fatalf("打开配置文件失败:%v", err)
}
monitor.InitConfig(&agentConfig)
if agentConfig.ClientSecret == "" {
log.Fatal("ClientSecret 不能为空")
}
if agentConfig.ReportDelay < 1 || agentConfig.ReportDelay > 4 {
log.Fatal("report-delay 的区间为 1-4")
}
if agentConfig.UUID == "" {
if uuid, err := uuid.GenerateUUID(); err == nil {
agentConfig.UUID = uuid
agentConfig.Save()
} else {
log.Fatalf("生成 UUID 失败:%v", err)
}
}
}
func run() {
auth := model.AuthHandler{
ClientSecret: agentConfig.ClientSecret,
@@ -645,8 +654,7 @@ func handleCommandTask(task *pb.Task, result *pb.TaskResult) {
return
}
startedAt := time.Now()
var cmd *exec.Cmd
var endCh = make(chan struct{})
endCh := make(chan struct{})
pg, err := processgroup.NewProcessExitGroup()
if err != nil {
// 进程组创建失败,直接退出
@@ -654,12 +662,14 @@ func handleCommandTask(task *pb.Task, result *pb.TaskResult) {
return
}
timeout := time.NewTimer(time.Hour * 2)
if util.IsWindows() {
cmd = exec.Command("cmd", "/c", task.GetData()) // #nosec
} else {
cmd = exec.Command("sh", "-c", task.GetData()) // #nosec
}
cmd := processgroup.NewCommand(task.GetData())
var b bytes.Buffer
cmd.Stdout = &b
cmd.Env = os.Environ()
if err = cmd.Start(); err != nil {
result.Data = err.Error()
return
}
pg.AddProcess(cmd)
go func() {
select {
@@ -671,12 +681,11 @@ func handleCommandTask(task *pb.Task, result *pb.TaskResult) {
timeout.Stop()
}
}()
output, err := cmd.Output()
if err != nil {
result.Data += fmt.Sprintf("%s\n%s", string(output), err.Error())
if err = cmd.Wait(); err != nil {
result.Data += fmt.Sprintf("%s\n%s", b.String(), err.Error())
} else {
close(endCh)
result.Data = string(output)
result.Data = b.String()
result.Successful = true
}
pg.Dispose()