diff --git a/internal/handler/agent_install_handler.go b/internal/handler/agent_install_handler.go index 4c5f6cf..10b570d 100644 --- a/internal/handler/agent_install_handler.go +++ b/internal/handler/agent_install_handler.go @@ -25,8 +25,21 @@ func (h *AgentHandler) DownloadAgent(c echo.Context) error { // 校验 API Key apiKey := c.QueryParam("key") if _, err := h.apiKeyService.ValidateApiKey(c.Request().Context(), apiKey); err != nil { - h.logger.Warn("download agent failed: invalid api key", zap.String("key", apiKey)) - return orz.NewError(401, "无效的 API 密钥") + // API Key 校验失败,尝试 IP 白名单兜底(兼容旧版 Agent 自动更新) + clientIP := c.RealIP() + isOnline, checkErr := h.agentService.IsAgentOnlineByIP(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, "系统内部错误") + } + + if !isOnline { + h.logger.Warn("download agent failed: invalid api key and ip not whitelisted", + zap.String("key", apiKey), + zap.String("ip", clientIP)) + return orz.NewError(401, "无效的 API 密钥,且来源 IP 未在白名单中") + } + h.logger.Info("download agent allowed by ip whitelist", zap.String("ip", clientIP), zap.String("filename", filename)) } // 从嵌入的文件系统读取 diff --git a/internal/service/agent_service.go b/internal/service/agent_service.go index 993b611..8d050a5 100644 --- a/internal/service/agent_service.go +++ b/internal/service/agent_service.go @@ -136,6 +136,20 @@ 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) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return false, nil + } + return false, err + } + + // 状态 1 表示在线 + return agent.Status == 1, nil +} + // HandleCommandResponse 处理指令响应 func (s *AgentService) HandleCommandResponse(ctx context.Context, agentID string, resp *protocol.CommandResponse) error { s.logger.Info("command response received",