feat: support disabling agents

internal/handler/agent_enabled_test.go

@@ -0,0 +1,81 @@
+package handler
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	"github.com/glebarez/sqlite"
+	"github.com/pika-monitor/pika/internal/models"
+	"github.com/pika-monitor/pika/internal/protocol"
+	"github.com/pika-monitor/pika/internal/repo"
+	"github.com/pika-monitor/pika/internal/service"
+	"go.uber.org/zap"
+	"gorm.io/gorm"
+)
+
+func TestHandleWebSocketMessageIgnoresDisabledAgent(t *testing.T) {
+	database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
+	if err != nil {
+		t.Fatalf("open database: %v", err)
+	}
+	if err := database.AutoMigrate(&models.Agent{}); err != nil {
+		t.Fatalf("migrate agent: %v", err)
+	}
+
+	agent := models.Agent{ID: "agent-disabled", Enabled: true}
+	if err := database.Create(&agent).Error; err != nil {
+		t.Fatalf("create agent: %v", err)
+	}
+	if err := database.Model(&models.Agent{}).Where("id = ?", agent.ID).Update("enabled", false).Error; err != nil {
+		t.Fatalf("disable agent: %v", err)
+	}
+
+	h := &AgentHandler{
+		logger: zap.NewNop(),
+		agentService: &service.AgentService{
+			AgentRepo: repo.NewAgentRepo(database),
+		},
+	}
+	invalidMetrics := json.RawMessage(`{`)
+	if err := h.handleWebSocketMessage(context.Background(), agent.ID, string(protocol.MessageTypeMetrics), invalidMetrics); err != nil {
+		t.Fatalf("disabled agent data should be ignored, got error: %v", err)
+	}
+
+	if err := database.Model(&models.Agent{}).Where("id = ?", agent.ID).Update("enabled", true).Error; err != nil {
+		t.Fatalf("enable agent: %v", err)
+	}
+	if err := h.handleWebSocketMessage(context.Background(), agent.ID, string(protocol.MessageTypeMetrics), invalidMetrics); err == nil {
+		t.Fatal("enabled agent data should reach the metrics parser")
+	}
+}
+
+func TestDisabledAgentStatusCannotBeRefreshed(t *testing.T) {
+	database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
+	if err != nil {
+		t.Fatalf("open database: %v", err)
+	}
+	if err := database.AutoMigrate(&models.Agent{}); err != nil {
+		t.Fatalf("migrate agent: %v", err)
+	}
+
+	agent := models.Agent{ID: "agent-status", Enabled: true, Status: 1}
+	if err := database.Create(&agent).Error; err != nil {
+		t.Fatalf("create agent: %v", err)
+	}
+	agentService := &service.AgentService{AgentRepo: repo.NewAgentRepo(database)}
+	if err := agentService.UpdateAgentEnabled(context.Background(), agent.ID, false); err != nil {
+		t.Fatalf("disable agent: %v", err)
+	}
+	if err := agentService.UpdateAgentStatus(context.Background(), agent.ID, 1); err != nil {
+		t.Fatalf("refresh disabled agent status: %v", err)
+	}
+
+	var saved models.Agent
+	if err := database.First(&saved, "id = ?", agent.ID).Error; err != nil {
+		t.Fatalf("load agent: %v", err)
+	}
+	if saved.Enabled || saved.Status != 0 {
+		t.Fatalf("disabled agent status changed unexpectedly: enabled=%v status=%d", saved.Enabled, saved.Status)
+	}
+}

internal/handler/agent_handler.go

@@ -1,13 +1,39 @@
 package handler
 
 import (
+	"sync"
+
+	"github.com/go-orz/orz"
 	"github.com/gorilla/websocket"
+	"github.com/labstack/echo/v5"
 	"github.com/pika-monitor/pika/internal/service"
 	ws "github.com/pika-monitor/pika/internal/websocket"
 	"go.uber.org/zap"
 )
 
+// Enable 启用探针数据处理。
+func (h *AgentHandler) Enable(c *echo.Context) error {
+	return h.updateEnabled(c, true)
+}
+
+// Disable 禁用探针数据处理和告警。
+func (h *AgentHandler) Disable(c *echo.Context) error {
+	return h.updateEnabled(c, false)
+}
+
+func (h *AgentHandler) updateEnabled(c *echo.Context, enabled bool) error {
+	h.enabledMu.Lock()
+	defer h.enabledMu.Unlock()
+
+	agentID := c.Param("id")
+	if err := h.agentService.UpdateAgentEnabled(c.Request().Context(), agentID, enabled); err != nil {
+		return err
+	}
+	return orz.Ok(c, orz.Map{})
+}
+
 type AgentHandler struct {
+	enabledMu       sync.RWMutex
 	logger          *zap.Logger
 	agentService    *service.AgentService
 	trafficService  *service.TrafficService

internal/handler/agent_query_handler.go

@@ -1,6 +1,8 @@
 package handler
 
 import (
+	"encoding/json"
+
 	"github.com/go-orz/orz"
 	"github.com/labstack/echo/v5"
 	"github.com/pika-monitor/pika/internal/models"
@@ -32,6 +34,19 @@ func (h *AgentHandler) Get(c *echo.Context) error {
 		agent.Hostname = ""
 	}
 
+	if !isAuthenticated {
+		payload, err := json.Marshal(agent)
+		if err != nil {
+			return err
+		}
+		var publicAgent map[string]any
+		if err := json.Unmarshal(payload, &publicAgent); err != nil {
+			return err
+		}
+		delete(publicAgent, "enabled")
+		return orz.Ok(c, publicAgent)
+	}
+
 	return orz.Ok(c, agent)
 }
 

internal/handler/agent_ws_handler.go

@@ -28,7 +28,9 @@ func (h *AgentHandler) HandleWebSocket(c *echo.Context) error {
 	}
 
 	// 注册探针 - 使用独立的context,不依赖HTTP请求的context
+	h.enabledMu.RLock()
 	agent, err := h.agentService.RegisterAgent(context.Background(), c.RealIP(), &registerReq.AgentInfo, registerReq.ApiKey)
+	h.enabledMu.RUnlock()
 	if err != nil {
 		// 发送注册失败响应
 		h.sendRegisterError(conn, err.Error())
@@ -47,20 +49,22 @@ func (h *AgentHandler) HandleWebSocket(c *echo.Context) error {
 		return err
 	}
 
-	// 下发防篡改配置
-	if err := h.sendTamperConfig(conn, agent.ID); err != nil {
-		h.logger.Error("failed to send tamper config", zap.Error(err))
-		// 配置下发失败不中断连接，只记录日志
-	}
-	// 下发SSH登录监控配置
-	if err := h.sendSSHLoginConfig(conn, agent.ID); err != nil {
-		h.logger.Error("failed to send ssh login config", zap.Error(err))
-		// 配置下发失败不中断连接，只记录日志
-	}
-	// 下发公网 IP 采集配置
-	if err := h.sendPublicIPConfig(conn, agent.ID); err != nil {
-		h.logger.Error("failed to send public ip config", zap.Error(err))
-		// 配置下发失败不中断连接，只记录日志
+	if agent.Enabled {
+		// 下发防篡改配置
+		if err := h.sendTamperConfig(conn, agent.ID); err != nil {
+			h.logger.Error("failed to send tamper config", zap.Error(err))
+			// 配置下发失败不中断连接，只记录日志
+		}
+		// 下发SSH登录监控配置
+		if err := h.sendSSHLoginConfig(conn, agent.ID); err != nil {
+			h.logger.Error("failed to send ssh login config", zap.Error(err))
+			// 配置下发失败不中断连接，只记录日志
+		}
+		// 下发公网 IP 采集配置
+		if err := h.sendPublicIPConfig(conn, agent.ID); err != nil {
+			h.logger.Error("failed to send public ip config", zap.Error(err))
+			// 配置下发失败不中断连接，只记录日志
+		}
 	}
 
 	// 创建客户端并注册到管理器
@@ -76,6 +80,17 @@ func (h *AgentHandler) HandleWebSocket(c *echo.Context) error {
 
 // handleWebSocketMessage 处理WebSocket消息
 func (h *AgentHandler) handleWebSocketMessage(ctx context.Context, agentID string, messageType string, data json.RawMessage) error {
+	h.enabledMu.RLock()
+	defer h.enabledMu.RUnlock()
+
+	enabled, err := h.agentService.IsAgentEnabled(ctx, agentID)
+	if err != nil {
+		return err
+	}
+	if !enabled {
+		return nil
+	}
+
 	switch protocol.MessageType(messageType) {
 	case protocol.MessageTypeMetrics:
 		return h.handleMetricsMessage(ctx, agentID, data)