From 81a08a38da8d55d94be3a0819e7db9a88f20a88f Mon Sep 17 00:00:00 2001
From: naiba <hi@nai.ba>
Date: Sat, 4 Jan 2025 01:10:16 +0800
Subject: [PATCH] improve: server list lock

---
 cmd/dashboard/controller/service.go |  4 +--
 service/rpc/nezha.go                | 42 ++++++++++++++++-------------
 2 files changed, 25 insertions(+), 21 deletions(-)

diff --git a/cmd/dashboard/controller/service.go b/cmd/dashboard/controller/service.go
index d1ac827..2a3022c 100644
--- a/cmd/dashboard/controller/service.go
+++ b/cmd/dashboard/controller/service.go
@@ -85,7 +85,8 @@ func listServiceHistory(c *gin.Context) ([]*model.ServiceInfos, error) {
 
 	singleton.ServerLock.RLock()
 	server, ok := singleton.ServerList[id]
-	if !ok {
+	singleton.ServerLock.RUnlock()
+	if !ok || server == nil {
 		return nil, singleton.Localizer.ErrorT("server not found")
 	}
 
@@ -95,7 +96,6 @@ func listServiceHistory(c *gin.Context) ([]*model.ServiceInfos, error) {
 	if server.HideForGuest && !authorized {
 		return nil, singleton.Localizer.ErrorT("unauthorized")
 	}
-	singleton.ServerLock.RUnlock()
 
 	var serviceHistories []*model.ServiceHistory
 	if err := singleton.DB.Model(&model.ServiceHistory{}).Select("service_id, created_at, server_id, avg_delay").
diff --git a/service/rpc/nezha.go b/service/rpc/nezha.go
index 8002f26..ba49661 100644
--- a/service/rpc/nezha.go
+++ b/service/rpc/nezha.go
@@ -103,20 +103,20 @@ func (s *NezhaHandler) ReportSystemState(stream pb.NezhaService_ReportSystemStat
 		state := model.PB2State(state)
 
 		singleton.ServerLock.RLock()
+		server, ok := singleton.ServerList[clientID]
+		singleton.ServerLock.RUnlock()
 
-		if singleton.ServerList[clientID] == nil {
-			singleton.ServerLock.RUnlock()
+		if !ok || server == nil {
 			return nil
 		}
 
-		singleton.ServerList[clientID].LastActive = time.Now()
-		singleton.ServerList[clientID].State = &state
+		server.LastActive = time.Now()
+		server.State = &state
 		// 应对 dashboard 重启的情况，如果从未记录过，先打点，等到小时时间点时入库
-		if singleton.ServerList[clientID].PrevTransferInSnapshot == 0 || singleton.ServerList[clientID].PrevTransferOutSnapshot == 0 {
-			singleton.ServerList[clientID].PrevTransferInSnapshot = int64(state.NetInTransfer)
-			singleton.ServerList[clientID].PrevTransferOutSnapshot = int64(state.NetOutTransfer)
+		if server.PrevTransferInSnapshot == 0 || server.PrevTransferOutSnapshot == 0 {
+			server.PrevTransferInSnapshot = int64(state.NetInTransfer)
+			server.PrevTransferOutSnapshot = int64(state.NetOutTransfer)
 		}
-		singleton.ServerLock.RUnlock()
 
 		stream.Send(&pb.Receipt{Proced: true})
 	}
@@ -129,6 +129,7 @@ func (s *NezhaHandler) onReportSystemInfo(c context.Context, r *pb.Host) error {
 		return err
 	}
 	host := model.PB2Host(r)
+
 	singleton.ServerLock.RLock()
 	defer singleton.ServerLock.RUnlock()
 
@@ -214,12 +215,18 @@ func (s *NezhaHandler) ReportGeoIP(c context.Context, r *pb.GeoIP) (*pb.GeoIP, e
 	joinedIP := geoip.IP.Join()
 
 	singleton.ServerLock.RLock()
+	server, ok := singleton.ServerList[clientID]
+	singleton.ServerLock.RUnlock()
+	if !ok || server == nil {
+		return nil, fmt.Errorf("server not found")
+	}
+
 	// 检查并更新DDNS
-	if singleton.ServerList[clientID].EnableDDNS && joinedIP != "" &&
-		(singleton.ServerList[clientID].GeoIP == nil || singleton.ServerList[clientID].GeoIP.IP != geoip.IP) {
+	if server.EnableDDNS && joinedIP != "" &&
+		(server.GeoIP == nil || server.GeoIP.IP != geoip.IP) {
 		ipv4 := geoip.IP.IPv4Addr
 		ipv6 := geoip.IP.IPv6Addr
-		providers, err := singleton.GetDDNSProvidersFromProfiles(singleton.ServerList[clientID].DDNSProfiles, &ddns.IP{Ipv4Addr: ipv4, Ipv6Addr: ipv6})
+		providers, err := singleton.GetDDNSProvidersFromProfiles(server.DDNSProfiles, &ddns.IP{Ipv4Addr: ipv4, Ipv6Addr: ipv6})
 		if err == nil {
 			for _, provider := range providers {
 				go func(provider *ddns.Provider) {
@@ -232,23 +239,22 @@ func (s *NezhaHandler) ReportGeoIP(c context.Context, r *pb.GeoIP) (*pb.GeoIP, e
 	}
 
 	// 发送IP变动通知
-	if singleton.ServerList[clientID].GeoIP != nil && singleton.Conf.EnableIPChangeNotification &&
+	if server.GeoIP != nil && singleton.Conf.EnableIPChangeNotification &&
 		((singleton.Conf.Cover == model.ConfigCoverAll && !singleton.Conf.IgnoredIPNotificationServerIDs[clientID]) ||
 			(singleton.Conf.Cover == model.ConfigCoverIgnoreAll && singleton.Conf.IgnoredIPNotificationServerIDs[clientID])) &&
-		singleton.ServerList[clientID].GeoIP.IP.Join() != "" &&
+		server.GeoIP.IP.Join() != "" &&
 		joinedIP != "" &&
-		singleton.ServerList[clientID].GeoIP.IP != geoip.IP {
+		server.GeoIP.IP != geoip.IP {
 
 		singleton.SendNotification(singleton.Conf.IPChangeNotificationGroupID,
 			fmt.Sprintf(
 				"[%s] %s, %s => %s",
 				singleton.Localizer.T("IP Changed"),
-				singleton.ServerList[clientID].Name, singleton.IPDesensitize(singleton.ServerList[clientID].GeoIP.IP.Join()),
+				server.Name, singleton.IPDesensitize(server.GeoIP.IP.Join()),
 				singleton.IPDesensitize(joinedIP),
 			),
 			nil)
 	}
-	singleton.ServerLock.RUnlock()
 
 	// 根据内置数据库查询 IP 地理位置
 	var ip string
@@ -266,9 +272,7 @@ func (s *NezhaHandler) ReportGeoIP(c context.Context, r *pb.GeoIP) (*pb.GeoIP, e
 	geoip.CountryCode = location
 
 	// 将地区码写入到 Host
-	singleton.ServerLock.Lock()
-	defer singleton.ServerLock.Unlock()
-	singleton.ServerList[clientID].GeoIP = &geoip
+	server.GeoIP = &geoip
 
 	return &pb.GeoIP{Ip: nil, CountryCode: location}, nil
 }