diff --git a/internal/handler/agent_enabled_test.go b/internal/handler/agent_enabled_test.go index 0978235..5209feb 100644 --- a/internal/handler/agent_enabled_test.go +++ b/internal/handler/agent_enabled_test.go @@ -107,4 +107,13 @@ func TestPublicAgentQueriesExcludeDisabledAgents(t *testing.T) { if len(publicAgents) != 1 || publicAgents[0].ID != "public-enabled" { t.Fatalf("unexpected public agents: %+v", publicAgents) } + + authenticatedAgents, err := (&service.AgentService{AgentRepo: repo.NewAgentRepo(database)}). + ListByAuth(context.Background(), true) + if err != nil { + t.Fatalf("query authenticated agents: %v", err) + } + if len(authenticatedAgents) != 1 || authenticatedAgents[0].ID != "public-enabled" { + t.Fatalf("authenticated query returned disabled agents: %+v", authenticatedAgents) + } } diff --git a/internal/repo/agent_repo.go b/internal/repo/agent_repo.go index b5bf5c1..33b781e 100644 --- a/internal/repo/agent_repo.go +++ b/internal/repo/agent_repo.go @@ -57,6 +57,15 @@ func (r *AgentRepo) FindOnlineAgents(ctx context.Context) ([]models.Agent, error return agents, err } +// FindEnabledAgents 查找所有启用的探针。 +func (r *AgentRepo) FindEnabledAgents(ctx context.Context) ([]models.Agent, error) { + var agents []models.Agent + err := r.db.WithContext(ctx). + Where("enabled = ?", true). + Find(&agents).Error + return agents, err +} + // FindByIP 根据公网IP查找探针 func (r *AgentRepo) FindByIP(ctx context.Context, ip string) (*models.Agent, error) { var agent models.Agent diff --git a/internal/service/agent_service.go b/internal/service/agent_service.go index 225b431..04cbf21 100644 --- a/internal/service/agent_service.go +++ b/internal/service/agent_service.go @@ -512,7 +512,7 @@ func (s *AgentService) DeleteAgent(ctx context.Context, agentID string) error { // ListByAuth 根据认证状态列出探针(已登录返回全部,未登录返回公开可见) func (s *AgentService) ListByAuth(ctx context.Context, isAuthenticated bool) ([]models.Agent, error) { if isAuthenticated { - return s.AgentRepo.FindAll(ctx) + return s.AgentRepo.FindEnabledAgents(ctx) } return s.AgentRepo.FindPublicAgents(ctx) } @@ -523,6 +523,9 @@ func (s *AgentService) GetAgentByAuth(ctx context.Context, id string, isAuthenti if isAuthenticated { agent, err := s.AgentRepo.FindById(ctx, id) if err == nil { + if !agent.Enabled { + return nil, gorm.ErrRecordNotFound + } return &agent, nil } if !errors.Is(err, gorm.ErrRecordNotFound) { @@ -533,6 +536,9 @@ func (s *AgentService) GetAgentByAuth(ctx context.Context, id string, isAuthenti if err != nil { return nil, err } + if !agent.Enabled { + return nil, gorm.ErrRecordNotFound + } return &agent, nil }