pika/internal/handler/agent_ws_handler.go

package handler

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
	"github.com/labstack/echo/v5"
	"github.com/pika-monitor/pika/internal/protocol"
	ws "github.com/pika-monitor/pika/internal/websocket"
	"go.uber.org/zap"
)

// HandleWebSocket 处理WebSocket连接
func (h *AgentHandler) HandleWebSocket(c *echo.Context) error {
	conn, err := h.upgrader.Upgrade(c.Response(), c.Request(), nil)
	if err != nil {
		h.logger.Error("failed to upgrade websocket", zap.Error(err))
		return err
	}

	registerReq, err := h.readRegisterRequest(conn)
	if err != nil {
		conn.Close()
		return err
	}

	// 注册探针 - 使用独立的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())
		conn.Close()
		return err
	}

	// 先原子激活连接及其 BootID 会话，再发送注册响应。这样响应中的
	// AckSeq 与当前连接属于同一次会话切换，旧连接无法在其间污染位点。
	client := ws.NewClient(agent.ID, conn, h.wsManager)
	ackSeq := h.wsManager.Register(client, registerReq.BootID)
	// RegisterAgent 的在线写入发生在连接切换之前，可能被恰好退出的旧连接
	// 覆盖；会话激活后再确认一次，确保数据库状态属于 current client。
	if err := h.agentService.UpdateAgentStatus(context.Background(), agent.ID, 1); err != nil {
		h.logger.Warn("failed to confirm agent online status", zap.String("agentID", agent.ID), zap.Error(err))
	}
	defer func() {
		if h.wsManager.Unregister(client) {
			h.markAgentOffline(client.ID)
		}
	}()

	// 发送注册成功响应：声明可靠投递支持，并回传该探针的累计确认
	// 位点，探针据此跳过已处理消息、只重放未确认部分
	if err := h.sendRegisterSuccess(conn, agent.ID, ackSeq); err != nil {
		h.logger.Error("failed to send register ack", zap.Error(err))
		conn.Close()
		return 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))
			// 配置下发失败不中断连接，只记录日志
		}
	}

	// 启动读写协程
	go client.WritePump()
	client.ReadPump(context.Background())
	return nil
}

// 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)

	case protocol.MessageTypeCommandResp:
		return h.handleCommandResponseMessage(ctx, agentID, data)

	case protocol.MessageTypeTamperEvent:
		return h.handleTamperEventMessage(ctx, agentID, data)

	case protocol.MessageTypeDDNSIPReport:
		return h.handleDDNSIPReportMessage(ctx, agentID, data)

	case protocol.MessageTypePublicIPReport:
		return h.handlePublicIPReportMessage(ctx, agentID, data)

	case protocol.MessageTypeSSHLoginEvent:
		return h.handleSSHLoginEventMessage(ctx, agentID, data)

	case protocol.MessageTypeSSHLoginConfigResult:
		return h.handleSSHLoginConfigResultMessage(ctx, agentID, data)

	case protocol.MessageTypeTamperProtect:
		return h.handleTamperProtectMessage(ctx, agentID, data)

	default:
		h.logger.Warn("unknown message type", zap.String("type", messageType))
		return nil
	}
}

func (h *AgentHandler) readRegisterRequest(conn *websocket.Conn) (*protocol.RegisterRequest, error) {
	conn.SetReadDeadline(time.Now().Add(10 * time.Second))
	_, message, err := conn.ReadMessage()
	if err != nil {
		h.logger.Error("failed to read register message", zap.Error(err))
		return nil, err
	}

	conn.SetReadDeadline(time.Time{})

	var msg protocol.InputMessage
	if err := json.Unmarshal(message, &msg); err != nil {
		h.logger.Error("failed to parse register message", zap.Error(err))
		return nil, err
	}

	if msg.Type != protocol.MessageTypeRegister {
		h.logger.Error("first message must be register", zap.String("type", string(msg.Type)))
		return nil, echo.NewHTTPError(http.StatusBadRequest, "首条消息必须是注册消息")
	}

	var registerReq protocol.RegisterRequest
	if err := json.Unmarshal(msg.Data, &registerReq); err != nil {
		h.logger.Error("failed to parse register request", zap.Error(err))
		return nil, err
	}

	return &registerReq, nil
}

func (h *AgentHandler) markAgentOffline(agentID string) {
	_ = h.agentService.UpdateAgentStatus(context.Background(), agentID, 0)
}

// pongStatusRefreshInterval 在线状态心跳刷新间隔：本连接首次 pong 立即
// 写库，之后按该间隔周期性刷新，避免每个 pong（10s/探针）都打数据库
const pongStatusRefreshInterval = 5 * time.Minute

func (h *AgentHandler) handleWebSocketPong(client *ws.Client) {
	if !client.ShouldWriteStatus(pongStatusRefreshInterval) {
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	current, err := h.wsManager.DoIfCurrent(client, func() error {
		return h.agentService.UpdateAgentStatus(ctx, client.ID, 1)
	})
	if !current {
		return
	}
	if err != nil {
		h.logger.Warn("failed to update agent status on pong", zap.String("agentID", client.ID), zap.Error(err))
		return
	}
	client.MarkStatusWritten()
}

func (h *AgentHandler) handleMetricsMessage(ctx context.Context, agentID string, data json.RawMessage) error {
	var batch protocol.MetricsBatch
	if err := json.Unmarshal(data, &batch); err != nil {
		return ws.Permanent(err)
	}

	var transientErr error
	var permanentErr error
	for _, sample := range batch.Samples {
		metricsData, err := json.Marshal(sample.Data)
		if err != nil {
			h.logger.Warn("failed to marshal metric sample", zap.Error(err))
			permanentErr = errors.Join(permanentErr, err)
			continue
		}
		if err := h.metricService.HandleMetricData(ctx, agentID, string(sample.Type), metricsData, sample.Timestamp); err != nil {
			h.logger.Warn("failed to handle metric sample", zap.Error(err), zap.String("type", string(sample.Type)))
			if isPayloadError(err) {
				permanentErr = errors.Join(permanentErr, err)
			} else {
				transientErr = errors.Join(transientErr, err)
			}
			continue
		}
	}
	if transientErr != nil {
		return transientErr
	}
	return ws.Permanent(permanentErr)
}

func (h *AgentHandler) handleCommandResponseMessage(ctx context.Context, agentID string, data json.RawMessage) error {
	var cmdResp protocol.CommandResponse
	if err := json.Unmarshal(data, &cmdResp); err != nil {
		return ws.Permanent(err)
	}
	return classifyMessageError(h.agentService.HandleCommandResponse(ctx, agentID, &cmdResp))
}

func isPayloadError(err error) bool {
	var syntaxErr *json.SyntaxError
	var typeErr *json.UnmarshalTypeError
	return errors.As(err, &syntaxErr) || errors.As(err, &typeErr)
}

func classifyMessageError(err error) error {
	if isPayloadError(err) {
		return ws.Permanent(err)
	}
	return err
}

func (h *AgentHandler) handleTamperEventMessage(ctx context.Context, agentID string, data json.RawMessage) error {
	var eventData protocol.TamperEventData
	if err := json.Unmarshal(data, &eventData); err != nil {
		h.logger.Error("failed to unmarshal tamper event", zap.Error(err))
		return ws.Permanent(err)
	}
	return h.tamperService.CreateEvent(ctx, agentID, &eventData)
}

func (h *AgentHandler) handleDDNSIPReportMessage(ctx context.Context, agentID string, data json.RawMessage) error {
	var ipReport protocol.DDNSIPReportData
	if err := json.Unmarshal(data, &ipReport); err != nil {
		h.logger.Error("failed to unmarshal ddns ip report", zap.Error(err))
		return ws.Permanent(err)
	}
	return h.ddnsService.HandleIPReport(ctx, agentID, &ipReport)
}

func (h *AgentHandler) handlePublicIPReportMessage(ctx context.Context, agentID string, data json.RawMessage) error {
	var ipReport protocol.PublicIPReportData
	if err := json.Unmarshal(data, &ipReport); err != nil {
		h.logger.Error("failed to unmarshal public ip report", zap.Error(err))
		return ws.Permanent(err)
	}
	return h.agentService.UpdatePublicIP(ctx, agentID, ipReport.IPv4, ipReport.IPv6)
}

func (h *AgentHandler) handleSSHLoginEventMessage(ctx context.Context, agentID string, data json.RawMessage) error {
	var eventData protocol.SSHLoginEvent
	if err := json.Unmarshal(data, &eventData); err != nil {
		h.logger.Error("failed to unmarshal ssh login event", zap.Error(err))
		return ws.Permanent(err)
	}
	return h.sshLoginService.HandleEvent(ctx, agentID, eventData)
}

func (h *AgentHandler) handleSSHLoginConfigResultMessage(ctx context.Context, agentID string, data json.RawMessage) error {
	var resultData protocol.SSHLoginConfigResult
	if err := json.Unmarshal(data, &resultData); err != nil {
		h.logger.Error("failed to unmarshal ssh login config result", zap.Error(err))
		return ws.Permanent(err)
	}
	return h.sshLoginService.HandleConfigResult(ctx, agentID, resultData)
}

func (h *AgentHandler) handleTamperProtectMessage(ctx context.Context, agentID string, data json.RawMessage) error {
	var protectResp protocol.TamperProtectResponse
	if err := json.Unmarshal(data, &protectResp); err != nil {
		h.logger.Error("failed to unmarshal tamper protect response", zap.Error(err))
		return ws.Permanent(err)
	}
	return h.tamperService.HandleConfigResult(ctx, agentID, protectResp)
}

// sendRegisterSuccess 发送注册成功响应
func (h *AgentHandler) sendRegisterSuccess(conn *websocket.Conn, agentID string, ackSeq uint64) error {
	resp := protocol.RegisterResponse{
		AgentID:  agentID,
		Status:   "success",
		Reliable: true,
		AckSeq:   ackSeq,
	}
	return conn.WriteJSON(protocol.OutboundMessage{
		Type: protocol.MessageTypeRegisterAck,
		Data: resp,
	})
}

// sendRegisterError 发送注册失败响应
func (h *AgentHandler) sendRegisterError(conn *websocket.Conn, errMsg string) error {
	resp := protocol.RegisterResponse{
		Status:  "error",
		Message: errMsg,
	}

	return conn.WriteJSON(protocol.OutboundMessage{
		Type: protocol.MessageTypeRegisterErr,
		Data: resp,
	})
}

// sendTamperConfig 发送防篡改配置（探针初始化时发送完整配置作为新增）
func (h *AgentHandler) sendTamperConfig(conn *websocket.Conn, agentID string) error {
	// 使用 TamperService 构建初始配置（复用逻辑，会自动判断 enabled 状态）
	added, removed, err := h.tamperService.BuildInitialConfig(context.Background(), agentID)
	if err != nil {
		return err
	}

	// 使用增量配置格式
	configData := protocol.TamperProtectConfig{
		Added:   added,
		Removed: removed,
	}

	msgData, err := json.Marshal(protocol.OutboundMessage{
		Type: protocol.MessageTypeTamperProtect,
		Data: configData,
	})
	if err != nil {
		return err
	}

	return conn.WriteMessage(websocket.TextMessage, msgData)
}

func (h *AgentHandler) sendSSHLoginConfig(conn *websocket.Conn, agentID string) error {
	config, err := h.sshLoginService.GetConfig(context.Background(), agentID)
	if err != nil {
		return err
	}
	msgData, err := json.Marshal(protocol.OutboundMessage{
		Type: protocol.MessageTypeSSHLoginConfig,
		Data: protocol.SSHLoginConfig{
			Enabled: config.Enabled,
		},
	})
	if err != nil {
		return err
	}
	return conn.WriteMessage(websocket.TextMessage, msgData)
}

func (h *AgentHandler) sendPublicIPConfig(conn *websocket.Conn, agentID string) error {
	config, err := h.propertyService.GetPublicIPConfig(context.Background())
	if err != nil {
		return err
	}

	if !config.Enabled || (!config.IPv4Enabled && !config.IPv6Enabled) {
		return nil
	}

	// 查询探针标签，用于按标签匹配采集范围
	var agentTags []string
	if agent, err := h.agentService.AgentRepo.FindById(context.Background(), agentID); err == nil {
		agentTags = agent.Tags
	}

	ipv4Enabled := config.IsIPv4Target(agentID, agentTags)
	ipv6Enabled := config.IsIPv6Target(agentID, agentTags)
	if !ipv4Enabled && !ipv6Enabled {
		return nil
	}

	msgData, err := json.Marshal(protocol.OutboundMessage{
		Type: protocol.MessageTypePublicIPConfig,
		Data: protocol.PublicIPConfigData{
			Enabled:         config.Enabled,
			IntervalSeconds: config.IntervalSeconds,
			IPv4Enabled:     ipv4Enabled,
			IPv6Enabled:     ipv6Enabled,
			IPv4APIs:        config.IPv4APIs,
			IPv6APIs:        config.IPv6APIs,
		},
	})
	if err != nil {
		return err
	}
	return conn.WriteMessage(websocket.TextMessage, msgData)
}