server/internal/handlers/agent_push.go

Rama main · réplica de lectura del repositorio original.

package handlers

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

	"blackbox/server/internal/hub"
	"blackbox/server/internal/middleware"
	"blackbox/server/internal/models"
	"blackbox/shared/types"
	"github.com/oklog/ulid/v2"
	"gorm.io/gorm"
)

const maxAgentEntryBodyBytes = 64 << 10

const (
	nodeStatusOnline    = "online"
	nodeStatusOffline   = "offline"
	nodeOfflineAfter    = 7 * time.Minute
	nodeStatusPollEvery = 30 * time.Second
)

func AgentPush(database *gorm.DB, h *hub.Hub, incidentCh chan<- types.Entry, shutdown <-chan struct{}) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		nodeName, ok := middleware.AgentNodeFromContext(r.Context())
		if !ok {
			writeError(w, http.StatusUnauthorized, "unauthorized")
			return
		}

		var entry types.Entry
		if !decodeJSONBody(w, r, maxAgentEntryBodyBytes, &entry) {
			return
		}
		if entry.ID == "" {
			writeError(w, http.StatusBadRequest, "entry id is required")
			return
		}
		if entry.NodeName != "" && entry.NodeName != nodeName {
			writeError(w, http.StatusForbidden, "agent node mismatch")
			return
		}
		entry.NodeName = nodeName
		serviceName := strings.ToLower(strings.TrimSpace(entry.Service))
		if serviceName == "" && !isAgentMetaEvent(entry) {
			writeError(w, http.StatusBadRequest, "service is required")
			return
		}
		entry.Service = serviceName
		if isExcluded(database, entry) {
			w.WriteHeader(http.StatusOK)
			return
		}
		if entry.Source == "docker" && entry.Event == "restart" {
			res := replaceDockerRestartEntry(database, entry, nodeName, h, incidentCh, shutdown)
			if res.Err != nil {
				writeError(w, http.StatusInternalServerError, "failed to process restart entry")
				return
			}
			if res.Found {
				if upsertNode(database, res.Updated) {
					broadcastNodeStatus(database, h)
				}
				w.WriteHeader(http.StatusOK)
				return
			}
		}
		created, err := createEntryIdempotent(database, entry, "agent-push")
		if errors.Is(err, errEntryIDConflict) {
			writeError(w, http.StatusConflict, "entry id already exists")
			return
		}
		if err != nil {
			writeError(w, http.StatusInternalServerError, "failed to save entry")
			return
		}
		if created {
			dispatchToIncidentChannelWithShutdown(incidentCh, shutdown, entry)
			if h != nil {
				if msg := MarshalWSMessage("entry", entry); msg != nil {
					h.Broadcast(msg)
				}
			}
		}
		if upsertNode(database, entry) {
			broadcastNodeStatus(database, h)
		}
		if created {
			w.WriteHeader(http.StatusCreated)
			return
		}
		w.WriteHeader(http.StatusOK)
	}
}

func isAgentMetaEvent(entry types.Entry) bool {
	return entry.Source == "agent" && (entry.Event == "heartbeat" || entry.Event == "start" || entry.Event == "shutdown")
}

func upsertNode(database *gorm.DB, entry types.Entry) bool {
	if entry.NodeName == "" {
		return false
	}

	now := time.Now().UTC()
	isMetaEvent := isAgentMetaEvent(entry)

	var node models.Node
	result := database.Where("name = ?", entry.NodeName).First(&node)

	if errors.Is(result.Error, gorm.ErrRecordNotFound) {
		newNode := models.Node{
			ID:       ulid.Make().String(),
			Name:     entry.NodeName,
			LastSeen: now,
			Status:   statusForEntry(entry),
		}
		if isMetaEvent {
			applyHeartbeatMeta(&newNode, entry.Metadata)
		}
		if err := database.Create(&newNode).Error; err != nil {
			log.Printf("upsertNode create error (node=%s): %v", entry.NodeName, err)
			return false
		}
		return true
	}

	if result.Error != nil {
		log.Printf("upsertNode lookup error (node=%s): %v", entry.NodeName, result.Error)
		return false
	}

	if !isMetaEvent && now.Sub(node.LastSeen) < 30*time.Second {
		return false
	}

	updates := map[string]interface{}{"last_seen": now}

	if isMetaEvent {
		updates["status"] = statusForEntry(entry)
		updatedNode := node
		applyHeartbeatMeta(&updatedNode, entry.Metadata)
		if updatedNode.AgentVersion != "" {
			updates["agent_version"] = updatedNode.AgentVersion
		}
		if updatedNode.IPAddress != "" {
			updates["ip_address"] = updatedNode.IPAddress
		}
		if updatedNode.OsInfo != "" {
			updates["os_info"] = updatedNode.OsInfo
		}
	}

	if err := database.Model(&node).Updates(updates).Error; err != nil {
		log.Printf("upsertNode update error (node=%s): %v", entry.NodeName, err)
		return false
	}
	return true
}

func applyHeartbeatMeta(node *models.Node, metadata string) {
	var meta struct {
		AgentVersion string `json:"agent_version"`
		IPAddress    string `json:"ip_address"`
		OsInfo       string `json:"os_info"`
	}
	if err := json.Unmarshal([]byte(metadata), &meta); err != nil {
		return
	}
	if meta.AgentVersion != "" {
		node.AgentVersion = meta.AgentVersion
	}
	if meta.IPAddress != "" {
		node.IPAddress = meta.IPAddress
	}
	if meta.OsInfo != "" {
		node.OsInfo = meta.OsInfo
	}
}

func broadcastNodeStatus(database *gorm.DB, h *hub.Hub) {
	if h == nil {
		return
	}
	var nodes []models.Node
	if err := database.Find(&nodes).Error; err != nil {
		return
	}
	online := 0
	now := time.Now()
	for _, n := range nodes {
		if effectiveNodeStatus(n, now) == nodeStatusOnline {
			online++
		}
	}
	type nodeStatus struct {
		Online int `json:"online"`
		Total  int `json:"total"`
	}
	if msg := MarshalWSMessage("node_status", nodeStatus{Online: online, Total: len(nodes)}); msg != nil {
		h.Broadcast(msg)
	}
}

func StartNodeStatusMonitor(ctx context.Context, database *gorm.DB, h *hub.Hub, interval time.Duration) {
	if ctx == nil || database == nil {
		return
	}
	if interval <= 0 {
		interval = nodeStatusPollEvery
	}

	go func() {
		ticker := time.NewTicker(interval)
		defer ticker.Stop()

		run := func() {
			changed, err := reconcileNodeStatuses(database, time.Now().UTC())
			if err != nil {
				log.Printf("node-status-monitor: reconcile failed: %v", err)
				return
			}
			if changed {
				broadcastNodeStatus(database, h)
			}
		}

		run()

		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				run()
			}
		}
	}()
}

func reconcileNodeStatuses(database *gorm.DB, now time.Time) (bool, error) {
	tx := database.Begin()
	if tx.Error != nil {
		return false, tx.Error
	}

	var nodes []models.Node
	if err := tx.Find(&nodes).Error; err != nil {
		tx.Rollback()
		return false, err
	}

	changed := false
	for _, node := range nodes {
		if !needsOfflineTransition(node, now) {
			continue
		}
		if err := tx.Model(&models.Node{}).Where("id = ?", node.ID).Update("status", nodeStatusOffline).Error; err != nil {
			tx.Rollback()
			return false, err
		}
		changed = true
	}

	if err := tx.Commit().Error; err != nil {
		tx.Rollback()
		return false, err
	}

	return changed, nil
}

func needsOfflineTransition(node models.Node, now time.Time) bool {
	return node.Status != nodeStatusOffline && effectiveNodeStatus(node, now) == nodeStatusOffline
}

func statusForEntry(entry types.Entry) string {
	if entry.Source == "agent" && entry.Event == "shutdown" {
		return nodeStatusOffline
	}
	return nodeStatusOnline
}

func effectiveNodeStatus(node models.Node, now time.Time) string {
	if node.Status == nodeStatusOffline {
		return nodeStatusOffline
	}
	if now.Sub(node.LastSeen) > nodeOfflineAfter {
		return nodeStatusOffline
	}
	return nodeStatusOnline
}