From 4d88d59100570e36372b9527800af184f9ead522 Mon Sep 17 00:00:00 2001
From: yusing <yusing@6uo.me>
Date: Tue, 15 Apr 2025 05:33:09 +0800
Subject: [PATCH] fix(agent): remove agent properly on verify error

---
 internal/config/query.go    | 9 ++++++---
 internal/utils/pool/pool.go | 4 ++++
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/internal/config/query.go b/internal/config/query.go
index bdc2668e..ec4cccf4 100644
--- a/internal/config/query.go
+++ b/internal/config/query.go
@@ -63,20 +63,23 @@ func (cfg *Config) VerifyNewAgent(host string, ca agent.PEMPair, client agent.PE
 		return 0, gperr.New("agent already exists")
 	}
 
-	var agentCfg agent.AgentConfig
+	agentCfg := new(agent.AgentConfig)
 	agentCfg.Addr = host
 	err := agentCfg.InitWithCerts(cfg.task.Context(), ca.Cert, client.Cert, client.Key)
 	if err != nil {
 		return 0, gperr.Wrap(err, "failed to start agent")
 	}
-	agent.Agents.Add(&agentCfg)
+	// must add it first to let LoadRoutes() reference from it
+	agent.Agents.Add(agentCfg)
 
-	provider := provider.NewAgentProvider(&agentCfg)
+	provider := provider.NewAgentProvider(agentCfg)
 	if err := cfg.errIfExists(provider); err != nil {
+		agent.Agents.Del(agentCfg)
 		return 0, err
 	}
 	err = provider.LoadRoutes()
 	if err != nil {
+		agent.Agents.Del(agentCfg)
 		return 0, gperr.Wrap(err, "failed to load routes")
 	}
 	return provider.NumRoutes(), nil
diff --git a/internal/utils/pool/pool.go b/internal/utils/pool/pool.go
index 9421ac44..69f67cda 100644
--- a/internal/utils/pool/pool.go
+++ b/internal/utils/pool/pool.go
@@ -26,6 +26,10 @@ func (p Pool[T]) Add(obj T) {
 	p.m.Store(obj.Key(), obj)
 }
 
+func (p Pool[T]) Del(obj T) {
+	p.m.Delete(obj.Key())
+}
+
 func (p Pool[T]) Get(key string) (T, bool) {
 	return p.m.Load(key)
 }