From c0288971d8b2e87c12173f184125849cd9b4cd8d Mon Sep 17 00:00:00 2001 From: UUBulb <35923940+uubulb@users.noreply.github.com> Date: Tue, 23 Jul 2024 23:20:55 +0800 Subject: [PATCH] chore: code style (#44) * some code refactor * service * chore: code style --------- Co-authored-by: naiba --- cmd/agent/main.go | 5 ++- cmd/agent/service.go | 1 + pkg/monitor/monitor.go | 79 +++++++++++++++++++++--------------------- pkg/monitor/myip.go | 31 ++++++----------- pkg/pty/pty_windows.go | 8 ++--- 5 files changed, 57 insertions(+), 67 deletions(-) diff --git a/cmd/agent/main.go b/cmd/agent/main.go index 9dfe4b96..180c067f 100644 --- a/cmd/agent/main.go +++ b/cmd/agent/main.go @@ -38,6 +38,7 @@ import ( pb "github.com/nezhahq/agent/proto" ) +// Agent 运行时参数。如需添加新参数，记得同时在 service.go 中添加 type AgentCliParam struct { SkipConnectionCount bool // 跳过连接数检查 SkipProcsCount bool // 跳过进程数量检查 @@ -432,7 +433,9 @@ func reportState() { // doSelfUpdate 执行更新检查 如果更新成功则会结束进程 func doSelfUpdate(useLocalVersion bool) { - <-monitor.Sync + if monitor.CachedCountry == "" { + return + } v := semver.MustParse("0.1.0") if useLocalVersion { v = semver.MustParse(version) diff --git a/cmd/agent/service.go b/cmd/agent/service.go index 87bfacc6..f01c21f0 100644 --- a/cmd/agent/service.go +++ b/cmd/agent/service.go @@ -86,6 +86,7 @@ func serviceActions(cmd *cobra.Command, args []string) { {agentCliParam.DisableForceUpdate, "--disable-force-update", ""}, {agentCliParam.UseIPv6CountryCode, "--use-ipv6-countrycode", ""}, {agentConfig.GPU, "--gpu", ""}, + {agentCliParam.UseGiteeToUpgrade, "--gitee", ""}, {agentCliParam.IPReportPeriod != 30*60, "-u", fmt.Sprint(agentCliParam.IPReportPeriod)}, } diff --git a/pkg/monitor/monitor.go b/pkg/monitor/monitor.go index cd13a244..97bf6d52 100644 --- a/pkg/monitor/monitor.go +++ b/pkg/monitor/monitor.go @@ -83,17 +83,14 @@ func GetHost() *model.Host { } else { if hi.VirtualizationRole == "guest" { cpuType = "Virtual" + ret.Virtualization = hi.VirtualizationSystem } else { cpuType = "Physical" + ret.Virtualization = "" } ret.Platform = hi.Platform ret.PlatformVersion = hi.PlatformVersion ret.Arch = hi.KernelArch - if cpuType == "Physical" { - ret.Virtualization = "" - } else { - ret.Virtualization = hi.VirtualizationSystem - } ret.BootTime = hi.BootTime } @@ -218,40 +215,6 @@ func GetState(skipConnectionCount bool, skipProcsCount bool) *model.HostState { } } - var tcpConnCount, udpConnCount uint64 - if !skipConnectionCount { - ss_err := true - if runtime.GOOS == "linux" { - tcpStat, err_tcp := goss.ConnectionsWithProtocol(goss.AF_INET, syscall.IPPROTO_TCP) - udpStat, err_udp := goss.ConnectionsWithProtocol(goss.AF_INET, syscall.IPPROTO_UDP) - if err_tcp == nil && err_udp == nil { - ss_err = false - tcpConnCount = uint64(len(tcpStat)) - udpConnCount = uint64(len(udpStat)) - } - if strings.Contains(CachedIP, ":") { - tcpStat6, err_tcp := goss.ConnectionsWithProtocol(goss.AF_INET6, syscall.IPPROTO_TCP) - udpStat6, err_udp := goss.ConnectionsWithProtocol(goss.AF_INET6, syscall.IPPROTO_UDP) - if err_tcp == nil && err_udp == nil { - ss_err = false - tcpConnCount += uint64(len(tcpStat6)) - udpConnCount += uint64(len(udpStat6)) - } - } - } - if ss_err { - conns, _ := net.Connections("all") - for i := 0; i < len(conns); i++ { - switch conns[i].Type { - case syscall.SOCK_STREAM: - tcpConnCount++ - case syscall.SOCK_DGRAM: - udpConnCount++ - } - } - } - } - if agentConfig.Temperature { go updateTemperatureStat() ret.Temperatures = temperatureStat @@ -265,7 +228,7 @@ func GetState(skipConnectionCount bool, skipProcsCount bool) *model.HostState { ret.NetInTransfer, ret.NetOutTransfer = netInTransfer, netOutTransfer ret.NetInSpeed, ret.NetOutSpeed = netInSpeed, netOutSpeed ret.Uptime = uint64(time.Since(cachedBootTime).Seconds()) - ret.TcpConnCount, ret.UdpConnCount = tcpConnCount, udpConnCount + ret.TcpConnCount, ret.UdpConnCount = getConns(skipConnectionCount) return &ret } @@ -352,6 +315,42 @@ func getDiskTotalAndUsed() (total uint64, used uint64) { return } +func getConns(skipConnectionCount bool) (tcpConnCount, udpConnCount uint64) { + if !skipConnectionCount { + ss_err := true + if runtime.GOOS == "linux" { + tcpStat, err_tcp := goss.ConnectionsWithProtocol(goss.AF_INET, syscall.IPPROTO_TCP) + udpStat, err_udp := goss.ConnectionsWithProtocol(goss.AF_INET, syscall.IPPROTO_UDP) + if err_tcp == nil && err_udp == nil { + ss_err = false + tcpConnCount = uint64(len(tcpStat)) + udpConnCount = uint64(len(udpStat)) + } + if strings.Contains(CachedIP, ":") { + tcpStat6, err_tcp := goss.ConnectionsWithProtocol(goss.AF_INET6, syscall.IPPROTO_TCP) + udpStat6, err_udp := goss.ConnectionsWithProtocol(goss.AF_INET6, syscall.IPPROTO_UDP) + if err_tcp == nil && err_udp == nil { + ss_err = false + tcpConnCount += uint64(len(tcpStat6)) + udpConnCount += uint64(len(udpStat6)) + } + } + } + if ss_err { + conns, _ := net.Connections("all") + for i := 0; i < len(conns); i++ { + switch conns[i].Type { + case syscall.SOCK_STREAM: + tcpConnCount++ + case syscall.SOCK_DGRAM: + udpConnCount++ + } + } + } + } + return tcpConnCount, udpConnCount +} + func updateGPUStat(gpuStat *uint64) { if !atomic.CompareAndSwapInt32(&updateGPUStatus, 0, 1) { return diff --git a/pkg/monitor/myip.go b/pkg/monitor/myip.go index 6094be51..3ede6aec 100644 --- a/pkg/monitor/myip.go +++ b/pkg/monitor/myip.go @@ -49,7 +49,6 @@ var ( // "https://freegeoip.app/json/", // 需要 Key } CachedIP, CachedCountry string - Sync = make(chan bool) httpClientV4 = util.NewSingleStackHTTPClient(time.Second*20, time.Second*5, time.Second*10, false) httpClientV6 = util.NewSingleStackHTTPClient(time.Second*20, time.Second*5, time.Second*10, true) ) @@ -58,16 +57,10 @@ var ( func UpdateIP(useIPv6CountryCode bool, period uint32) { for { util.Println(agentConfig.Debug, "正在更新本地缓存IP信息") - var primaryIP, secondaryIP geoIP - if useIPv6CountryCode { - primaryIP = fetchGeoIP(geoIPApiList, true) - secondaryIP = fetchGeoIP(geoIPApiList, false) - } else { - primaryIP = fetchGeoIP(geoIPApiList, false) - secondaryIP = fetchGeoIP(geoIPApiList, true) - } + ipv4 := fetchGeoIP(geoIPApiList, false) + ipv6 := fetchGeoIP(geoIPApiList, true) - if primaryIP.IP == "" && secondaryIP.IP == "" { + if ipv4.IP == "" && ipv6.IP == "" { if period > 60 { time.Sleep(time.Minute) } else { @@ -75,21 +68,17 @@ func UpdateIP(useIPv6CountryCode bool, period uint32) { } continue } - if primaryIP.IP == "" || secondaryIP.IP == "" { - CachedIP = fmt.Sprintf("%s%s", primaryIP.IP, secondaryIP.IP) + if ipv4.IP == "" || ipv6.IP == "" { + CachedIP = fmt.Sprintf("%s%s", ipv4.IP, ipv6.IP) } else { - CachedIP = fmt.Sprintf("%s/%s", primaryIP.IP, secondaryIP.IP) + CachedIP = fmt.Sprintf("%s/%s", ipv4.IP, ipv6.IP) } - if primaryIP.CountryCode != "" { - CachedCountry = primaryIP.CountryCode - } else if secondaryIP.CountryCode != "" { - CachedCountry = secondaryIP.CountryCode + if ipv4.CountryCode != "" { + CachedCountry = ipv4.CountryCode } - - select { - case Sync <- true: - default: + if ipv6.CountryCode != "" && (useIPv6CountryCode || CachedCountry == "") { + CachedCountry = ipv6.CountryCode } time.Sleep(time.Second * time.Duration(period)) diff --git a/pkg/pty/pty_windows.go b/pkg/pty/pty_windows.go index 65152287..97fb1a03 100644 --- a/pkg/pty/pty_windows.go +++ b/pkg/pty/pty_windows.go @@ -58,9 +58,8 @@ func VersionCheck() bool { } return version >= 17763 - } else { - return false } + return false } func DownloadDependency() { @@ -131,10 +130,9 @@ func Start() (Pty, error) { if !isWin10 { tty, err := winpty.OpenDefault(path, shellPath) return &winPTY{tty: tty}, err - } else { - tty, err := conpty.Start(shellPath, conpty.ConPtyWorkDir(path)) - return &conPty{tty: tty}, err } + tty, err := conpty.Start(shellPath, conpty.ConPtyWorkDir(path)) + return &conPty{tty: tty}, err } func (w *winPTY) Write(p []byte) (n int, err error) {