Skip to content

Commit

Permalink
refactor
Browse files Browse the repository at this point in the history
  • Loading branch information
uubulb committed Jan 30, 2025
1 parent d677f75 commit 2fb62c2
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 45 deletions.
29 changes: 19 additions & 10 deletions cmd/agent/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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
}

Expand All @@ -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{}{}
})
}
Expand Down
49 changes: 14 additions & 35 deletions model/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ func (c *AgentConfig) Read(path string) error {
}
}

return ValidateConfig(c)
return ValidateConfig(c, false)
}

func (c *AgentConfig) Save() error {
Expand All @@ -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
Expand All @@ -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
}
Expand All @@ -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
Expand Down

0 comments on commit 2fb62c2

Please sign in to comment.