Netpulse_SasS/server/internal/store/agents.go

package store

import (
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"errors"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"
	npv1 "github.com/netpulse/netpulse/gen/go/netpulse/v1"
)

var ErrAgentNotFound = errors.New("зонд не знайдено або токен відкликано")

// Agent — ідентичність зонда після автентифікації.
type Agent struct {
	ID       string
	TenantID string
	Name     string
	SiteID   string
	Modules  []string
	Limits   Limits
}

// Limits — ліміти опитування, які сервер диктує зонду в Welcome.
type Limits struct {
	MaxConcurrency   int
	IcmpRatePPS      int
	BatchSize        int
	BatchInterval    time.Duration
	MaxInFlight      int
	HeartbeatEvery   time.Duration
}

func defaultLimits() Limits {
	return Limits{
		MaxConcurrency: 256,
		IcmpRatePPS:    500,
		BatchSize:      500,
		BatchInterval:  5 * time.Second,
		MaxInFlight:    4,
		HeartbeatEvery: 30 * time.Second,
	}
}

// AuthenticateAgent знаходить зонда за токеном.
//
// У БД лежить лише sha256 — токен у відкритому вигляді не зберігається
// ніде. Порівняння хешів робиться в SQL за унікальним індексом; на
// знайденому рядку додатково звіряємо constant-time, щоб форма запиту
// не залежала від того, як індекс порівнював байти.
func (s *Store) AuthenticateAgent(ctx context.Context, token string) (*Agent, error) {
	if token == "" {
		return nil, ErrAgentNotFound
	}
	sum := sha256.Sum256([]byte(token))

	var (
		a       Agent
		siteID  *string
		hash    []byte
		status  string
		limits  map[string]any
	)

	err := s.pool.QueryRow(ctx, `
		SELECT id::text, tenant_id::text, name, site_id::text, token_hash,
		       status::text, enabled_modules, limits
		FROM core.agents
		WHERE token_hash = $1
	`, sum[:]).Scan(&a.ID, &a.TenantID, &a.Name, &siteID, &hash, &status, &a.Modules, &limits)

	if errors.Is(err, pgx.ErrNoRows) {
		return nil, ErrAgentNotFound
	}
	if err != nil {
		return nil, err
	}
	if subtle.ConstantTimeCompare(hash, sum[:]) != 1 {
		return nil, ErrAgentNotFound
	}
	if status == "disabled" {
		return nil, fmt.Errorf("%w: зонд вимкнено", ErrAgentNotFound)
	}
	if siteID != nil {
		a.SiteID = *siteID
	}

	a.Limits = defaultLimits()
	applyLimitOverrides(&a.Limits, limits)
	return &a, nil
}

func applyLimitOverrides(l *Limits, m map[string]any) {
	num := func(key string) (int, bool) {
		v, ok := m[key]
		if !ok {
			return 0, false
		}
		switch n := v.(type) {
		case float64:
			return int(n), true
		case int64:
			return int(n), true
		}
		return 0, false
	}
	if v, ok := num("max_concurrency"); ok && v > 0 {
		l.MaxConcurrency = v
	}
	if v, ok := num("icmp_rate_pps"); ok && v > 0 {
		l.IcmpRatePPS = v
	}
	if v, ok := num("batch_size"); ok && v > 0 {
		l.BatchSize = v
	}
	if v, ok := num("max_in_flight"); ok && v > 0 {
		l.MaxInFlight = v
	}
}

// MarkAgentOnline фіксує підключення зонда та його версію.
func (s *Store) MarkAgentOnline(ctx context.Context, a *Agent, hello *npv1.Hello) error {
	build := hello.GetBuild()
	return s.InTenantTx(ctx, a.TenantID, func(tx pgx.Tx) error {
		_, err := tx.Exec(ctx, `
			UPDATE core.agents
			SET status = 'online',
			    last_heartbeat_at = now(),
			    version = COALESCE(NULLIF($2,''), version),
			    os      = COALESCE(NULLIF($3,''), os),
			    arch    = COALESCE(NULLIF($4,''), arch),
			    hostname= COALESCE(NULLIF($5,''), hostname),
			    updated_at = now()
			WHERE id = $1 AND tenant_id = $6
		`, a.ID, build.GetVersion(), build.GetOs(), build.GetArch(),
			hello.GetHostname(), a.TenantID)
		return err
	})
}

// MarkAgentOffline викликається при розриві сесії.
func (s *Store) MarkAgentOffline(ctx context.Context, a *Agent) error {
	return s.InTenantTx(ctx, a.TenantID, func(tx pgx.Tx) error {
		_, err := tx.Exec(ctx, `
			UPDATE core.agents SET status = 'offline', updated_at = now()
			WHERE id = $1 AND tenant_id = $2
		`, a.ID, a.TenantID)
		return err
	})
}

// RecordHeartbeat пише самометрики зонда в ts.agent_health і оновлює
// зведення в core.agents.health для швидкого показу в UI.
func (s *Store) RecordHeartbeat(ctx context.Context, a *Agent, hb *npv1.Heartbeat) error {
	h := hb.GetHealth()
	ts := hb.GetTs().AsTime()
	if ts.IsZero() {
		ts = time.Now()
	}

	batch := &pgx.Batch{}
	batch.Queue(`
		INSERT INTO ts.agent_health
			(ts, agent_id, tenant_id, cpu_pct, rss_bytes, goroutines,
			 queue_depth, checks_per_sec, errors_per_min)
		VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9)
		ON CONFLICT (ts, agent_id) DO NOTHING
	`, ts, a.ID, a.TenantID, h.GetCpuPct(), int64(h.GetRssBytes()),
		int32(h.GetGoroutines()), int32(h.GetQueueDepth()),
		h.GetChecksPerSec(), h.GetErrorsPerMin())

	// dropped_samples у зведенні — це видима діра в даних; вона має
	// бути помітна оператору, а не тільки в графіку самометрик.
	batch.Queue(`
		UPDATE core.agents
		SET last_heartbeat_at = now(),
		    status = 'online',
		    health = jsonb_build_object(
		        'rss_bytes', $3::bigint,
		        'queue_depth', $4::int,
		        'dropped_samples', $5::bigint,
		        'tasks_running', $6::int,
		        'clock_skew_ms', $7::bigint
		    )
		WHERE id = $1 AND tenant_id = $2
	`, a.ID, a.TenantID, int64(h.GetRssBytes()), int32(h.GetQueueDepth()),
		int64(h.GetDroppedSamples()), int32(hb.GetTasksRunning()),
		h.GetClockSkew().AsDuration().Milliseconds())

	res := s.pool.SendBatch(ctx, batch)
	defer res.Close()
	for i := 0; i < batch.Len(); i++ {
		if _, err := res.Exec(); err != nil {
			return fmt.Errorf("heartbeat[%d]: %w", i, err)
		}
	}
	return nil
}

// RecordTaskStatus оновлює core.checks за доповіддю агента.
func (s *Store) RecordTaskStatus(ctx context.Context, a *Agent, u *npv1.TaskStatusUpdate) error {
	var errText any
	if u.GetError() != nil {
		errText = u.GetError().GetCode() + ": " + u.GetError().GetMessage()
	}

	return s.InTenantTx(ctx, a.TenantID, func(tx pgx.Tx) error {
		_, err := tx.Exec(ctx, `
			UPDATE core.checks
			SET last_run_at = COALESCE($3, now()),
			    last_error  = $4,
			    updated_at  = now()
			WHERE id = $1 AND tenant_id = $2
		`, u.GetCheckId(), a.TenantID, tsOrNil(u), errText)
		return err
	})
}

func tsOrNil(u *npv1.TaskStatusUpdate) any {
	if u.GetTs() == nil {
		return nil
	}
	return u.GetTs().AsTime()
}