diff --git a/internal/handler/agent_install_handler.go b/internal/handler/agent_install_handler.go index 021e796..d9ddab5 100644 --- a/internal/handler/agent_install_handler.go +++ b/internal/handler/agent_install_handler.go @@ -27,7 +27,7 @@ func (h *AgentHandler) DownloadAgent(c echo.Context) error { if _, err := h.apiKeyService.ValidateApiKey(c.Request().Context(), apiKey); err != nil { // API Key 校验失败,尝试 IP 白名单兜底(兼容旧版 Agent 自动更新) clientIP := c.RealIP() - isOnline, checkErr := h.agentService.IsAgentOnlineByIP(c.Request().Context(), clientIP) + isOnline, checkErr := h.agentService.IsAgentByIP(c.Request().Context(), clientIP) if checkErr != nil { h.logger.Error("check agent online by ip failed", zap.String("ip", clientIP), zap.Error(checkErr)) return orz.NewError(500, "系统内部错误") diff --git a/internal/models/agent.go b/internal/models/agent.go index e1539d3..b2c7def 100644 --- a/internal/models/agent.go +++ b/internal/models/agent.go @@ -12,6 +12,7 @@ type Agent struct { ID string `gorm:"primaryKey" json:"id"` // 探针ID (UUID) Name string `gorm:"index" json:"name"` // 探针名称 Hostname string `gorm:"index" json:"hostname,omitempty"` // 主机名 + IP string `gorm:"index" json:"ip,omitempty"` // 连接 IP 地址 IPv4 string `gorm:"index" json:"ipv4,omitempty"` // 公网 IPv4 地址 IPv6 string `gorm:"index" json:"ipv6,omitempty"` // 公网 IPv6 地址 OS string `json:"os"` // 操作系统 diff --git a/internal/repo/agent_repo.go b/internal/repo/agent_repo.go index 6ad3a61..d227b9d 100644 --- a/internal/repo/agent_repo.go +++ b/internal/repo/agent_repo.go @@ -48,7 +48,7 @@ func (r *AgentRepo) FindOnlineAgents(ctx context.Context) ([]models.Agent, error func (r *AgentRepo) FindByIP(ctx context.Context, ip string) (*models.Agent, error) { var agent models.Agent err := r.db.WithContext(ctx). - Where("ipv4 = ? OR ipv6 = ?", ip, ip). + Where("ip = ? or ipv4 = ? OR ipv6 = ?", ip, ip). First(&agent).Error if err != nil { return nil, err diff --git a/internal/service/agent_service.go b/internal/service/agent_service.go index 8cab90a..e8be6bf 100644 --- a/internal/service/agent_service.go +++ b/internal/service/agent_service.go @@ -63,6 +63,7 @@ func (s *AgentService) RegisterAgent(ctx context.Context, ip string, info *proto // 更新现有探针信息(允许主机名、名称等变化) now := time.Now().UnixMilli() existingAgent.Hostname = info.Hostname + existingAgent.IP = ip existingAgent.OS = info.OS existingAgent.Arch = info.Arch existingAgent.Version = info.Version @@ -88,6 +89,7 @@ func (s *AgentService) RegisterAgent(ctx context.Context, ip string, info *proto ID: info.ID, // 使用客户端持久化的 ID Name: info.Name, Hostname: info.Hostname, + IP: ip, OS: info.OS, Arch: info.Arch, Version: info.Version, @@ -151,18 +153,16 @@ func (s *AgentService) ListOnlineAgents(ctx context.Context) ([]models.Agent, er return s.AgentRepo.FindOnlineAgents(ctx) } -// IsAgentOnlineByIP 检查指定公网IP是否存在在线的探针 -func (s *AgentService) IsAgentOnlineByIP(ctx context.Context, ip string) (bool, error) { - agent, err := s.AgentRepo.FindByIP(ctx, ip) +// IsAgentByIP 检查指定公网IP是否为探针 +func (s *AgentService) IsAgentByIP(ctx context.Context, ip string) (bool, error) { + _, err := s.AgentRepo.FindByIP(ctx, ip) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return false, nil } return false, err } - - // 状态 1 表示在线 - return agent.Status == 1, nil + return true, nil } // HandleCommandResponse 处理指令响应