chore: code style (#44)

* some code refactor * service * chore: code style --------- Co-authored-by: naiba <[email protected]>
nezhahq · Jul 23, 2024 · c028897 · c028897
1 parent a99bf70
commit c028897
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 67 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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,38 +57,28 @@ 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 {
 				time.Sleep(time.Second * time.Duration(period))
 			}
 			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
@@ -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) {