From 3a8caec66e3619821d67af54aa8c01503fb4e171 Mon Sep 17 00:00:00 2001 From: dushixiang Date: Sat, 17 Jan 2026 17:14:19 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=9D=E7=95=99=E6=97=A7=E7=9A=84=E8=BF=9E?= =?UTF-8?q?=E6=8E=A5=20IP=20=E5=9C=B0=E5=9D=80?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/handler/agent_install_handler.go | 2 +- internal/models/agent.go | 1 + internal/repo/agent_repo.go | 2 +- internal/service/agent_service.go | 12 ++++++------ 4 files changed, 9 insertions(+), 8 deletions(-) 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 处理指令响应