From 2fb62c2ce9fd2fdc9a6f97ef88eceeb9163a2ce8 Mon Sep 17 00:00:00 2001
From: uubulb <uub@kuzu.uk>
Date: Thu, 30 Jan 2025 21:57:59 +0800
Subject: [PATCH] refactor

---
 cmd/agent/main.go | 29 ++++++++++++++++++----------
 model/config.go   | 49 ++++++++++++++---------------------------------
 2 files changed, 33 insertions(+), 45 deletions(-)

diff --git a/cmd/agent/main.go b/cmd/agent/main.go
index 6d44746..66296a9 100644
--- a/cmd/agent/main.go
+++ b/cmd/agent/main.go
@@ -59,8 +59,9 @@ var (
 	lastReportHostInfo    time.Time
 	lastReportIPInfo      time.Time
 
-	hostStatus atomic.Bool
-	ipStatus   atomic.Bool
+	hostStatus   atomic.Bool
+	ipStatus     atomic.Bool
+	reloadStatus atomic.Bool
 
 	dnsResolver = &net.Resolver{PreferGo: true}
 	httpClient  = &http.Client{
@@ -470,7 +471,7 @@ func doTask(task *pb.Task) *pb.TaskResult {
 	case model.TaskTypeReportConfig:
 		handleReportConfigTask(&result)
 	case model.TaskTypeApplyConfig:
-		handleApplyConfigTask(task, &result)
+		handleApplyConfigTask(task)
 	case model.TaskTypeKeepalive:
 	default:
 		printf("不支持的任务: %v", task)
@@ -819,6 +820,11 @@ func handleReportConfigTask(result *pb.TaskResult) {
 		return
 	}
 
+	if reloadStatus.Load() {
+		result.Data = "another reload is in process"
+		return
+	}
+
 	println("Executing Report Config Task")
 
 	c, err := util.Json.Marshal(agentConfig)
@@ -831,9 +837,12 @@ func handleReportConfigTask(result *pb.TaskResult) {
 	result.Successful = true
 }
 
-func handleApplyConfigTask(task *pb.Task, result *pb.TaskResult) {
+func handleApplyConfigTask(task *pb.Task) {
 	if agentConfig.DisableCommandExecute {
-		result.Data = "此 Agent 已禁止命令执行"
+		return
+	}
+
+	if !reloadStatus.CompareAndSwap(false, true) {
 		return
 	}
 
@@ -842,25 +851,25 @@ func handleApplyConfigTask(task *pb.Task, result *pb.TaskResult) {
 	var tmpConfig model.AgentConfig
 	json := []byte(task.GetData())
 	if err := util.Json.Unmarshal(json, &tmpConfig); err != nil {
-		result.Data = err.Error()
+		printf("Validate Config failed: %v", err)
 		return
 	}
 
-	if err := model.ValidateConfig(&tmpConfig); err != nil {
-		result.Data = err.Error()
+	if err := model.ValidateConfig(&tmpConfig, true); err != nil {
+		printf("Validate Config failed: %v", err)
 		return
 	}
 
-	result.Successful = true
 	println("Will reload workers in 30 seconds")
 	time.AfterFunc(time.Second*30, func() {
 		println("Applying new configuration...")
 		agentConfig.Apply(&tmpConfig)
-		agentConfig.SaveToYAML(json)
+		agentConfig.Save()
 		geoipReported = false
 		logger.SetEnable(agentConfig.Debug)
 		monitor.InitConfig(&agentConfig)
 		monitor.CustomEndpoints = agentConfig.CustomIPApi
+		reloadStatus.Store(false)
 		reloadSigChan <- struct{}{}
 	})
 }
diff --git a/model/config.go b/model/config.go
index a00b622..cbc56bb 100644
--- a/model/config.go
+++ b/model/config.go
@@ -84,7 +84,7 @@ func (c *AgentConfig) Read(path string) error {
 		}
 	}
 
-	return ValidateConfig(c)
+	return ValidateConfig(c, false)
 }
 
 func (c *AgentConfig) Save() error {
@@ -101,28 +101,10 @@ func (c *AgentConfig) Save() error {
 	return os.WriteFile(c.filePath, data, 0600)
 }
 
-func (c *AgentConfig) SaveToYAML(j []byte) error {
-	data, err := yaml.JSONToYAML(j)
-	if err != nil {
-		return err
-	}
-
-	dir := filepath.Dir(c.filePath)
-	if err := os.MkdirAll(dir, 0750); err != nil {
-		return err
-	}
-
-	return os.WriteFile(c.filePath, data, 0600)
-}
-
 func (c *AgentConfig) Apply(new *AgentConfig) {
 	c.Debug = new.Debug
-	c.Server = new.Server
-	c.ClientSecret = new.ClientSecret
-	c.UUID = new.UUID
 	c.HardDrivePartitionAllowlist = new.HardDrivePartitionAllowlist
 	c.NICAllowlist = new.NICAllowlist
-	c.DNS = new.DNS
 	c.GPU = new.GPU
 	c.Temperature = new.Temperature
 	c.SkipConnectionCount = new.SkipConnectionCount
@@ -131,18 +113,13 @@ func (c *AgentConfig) Apply(new *AgentConfig) {
 	c.DisableForceUpdate = new.DisableForceUpdate
 	c.DisableCommandExecute = new.DisableCommandExecute
 	c.ReportDelay = new.ReportDelay
-	c.TLS = new.TLS
-	c.InsecureTLS = new.InsecureTLS
-	c.UseIPv6CountryCode = new.UseIPv6CountryCode
-	c.UseGiteeToUpgrade = new.UseGiteeToUpgrade
 	c.DisableNat = new.DisableNat
 	c.DisableSendQuery = new.DisableSendQuery
 	c.IPReportPeriod = new.IPReportPeriod
 	c.SelfUpdatePeriod = new.SelfUpdatePeriod
-	c.CustomIPApi = new.CustomIPApi
 }
 
-func ValidateConfig(c *AgentConfig) error {
+func ValidateConfig(c *AgentConfig, isRemoteEdit bool) error {
 	if c.ReportDelay == 0 {
 		c.ReportDelay = 3
 	}
@@ -153,20 +130,22 @@ func ValidateConfig(c *AgentConfig) error {
 		c.IPReportPeriod = 30
 	}
 
-	if c.Server == "" {
-		return errors.New("server address should not be empty")
-	}
-
-	if c.ClientSecret == "" {
-		return errors.New("client_secret must be specified")
-	}
-
 	if c.ReportDelay < 1 || c.ReportDelay > 4 {
 		return errors.New("report-delay ranges from 1-4")
 	}
 
-	if _, err := uuid.ParseUUID(c.UUID); err != nil {
-		return err
+	if !isRemoteEdit {
+		if c.Server == "" {
+			return errors.New("server address should not be empty")
+		}
+
+		if c.ClientSecret == "" {
+			return errors.New("client_secret must be specified")
+		}
+
+		if _, err := uuid.ParseUUID(c.UUID); err != nil {
+			return err
+		}
 	}
 
 	return nil