lpra/internal/server/handlers_panel_stat.go

package server

import (
	"crypto/rand"
	"errors"
	"io"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/x-grk/longpolling-remote-access/internal/proto"
)

const (
	// statTaskTimeout bounds how long a status request waits for the agent.
	statTaskTimeout = 8 * time.Second
	// statCacheFor is how long one sample answers every open status page, so
	// several tabs cost the agent a single reading.
	statCacheFor = 2 * time.Second
	// speedTaskTimeout covers the agent's ping + down + up windows with slack.
	speedTaskTimeout = 30 * time.Second
	// speedDownCap bounds one download stream; speedDownFor ends it if the
	// agent does not hang up first.
	speedDownCap = 256 << 20
	speedDownFor = 6 * time.Second
	speedUpCap   = 64 << 20
)

// statCache holds the last sample per agent and serializes its refresh.
type statCache struct {
	mu sync.Mutex
	m  map[string]*statEntry
}

type statEntry struct {
	mu  sync.Mutex
	at  time.Time
	res proto.TaskResult
}

func (c *statCache) entry(agentID string) *statEntry {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.m == nil {
		c.m = map[string]*statEntry{}
	}
	e := c.m[agentID]
	if e == nil {
		e = &statEntry{}
		c.m[agentID] = e
	}
	return e
}

func (c *statCache) drop(agentID string) {
	c.mu.Lock()
	delete(c.m, agentID)
	c.mu.Unlock()
}

func (s *Server) handleStatusPage(w http.ResponseWriter, r *http.Request) {
	a, err := s.db.GetAgent(r.PathValue("id"))
	if err != nil || a == nil {
		http.NotFound(w, r)
		return
	}
	s.render(w, "status.html", map[string]any{"Agent": a})
}

// runTask runs a one-shot task on the agent and waits for its result.
func (s *Server) runTask(agentID, typ string, timeout time.Duration) (proto.TaskResult, error) {
	taskID, ch := s.hub.CreateTask(agentID, typ, nil)
	defer s.hub.DiscardTask(taskID)
	t := time.NewTimer(timeout)
	defer t.Stop()
	select {
	case res := <-ch:
		return res, nil
	case <-t.C:
		return proto.TaskResult{}, errors.New("таймаут ожидания клиента")
	}
}

// onlineAgent answers the request itself and returns false when the agent is
// unknown or offline.
func (s *Server) onlineAgent(w http.ResponseWriter, r *http.Request) (string, bool) {
	id := r.PathValue("id")
	a, err := s.db.GetAgent(id)
	if err != nil || a == nil {
		http.NotFound(w, r)
		return "", false
	}
	if !s.online(a) {
		writeJSON(w, proto.TaskResult{Error: "машина офлайн"})
		return "", false
	}
	return id, true
}

// oldAgentError replaces the "unknown task" answer of an agent that predates
// the task.
func oldAgentError(res proto.TaskResult) proto.TaskResult {
	if !res.OK && strings.Contains(res.Error, "неизвестная операция") {
		res.Error = "клиент не поддерживает эту функцию - обновите его (нужен 0.0.5+)"
	}
	return res
}

// GET /api/panel/agents/{id}/stat - CPU, RAM and disk of the machine right now.
func (s *Server) handleAgentStat(w http.ResponseWriter, r *http.Request) {
	id, ok := s.onlineAgent(w, r)
	if !ok {
		return
	}
	e := s.stats.entry(id)
	e.mu.Lock()
	defer e.mu.Unlock()
	if time.Since(e.at) > statCacheFor {
		res, err := s.runTask(id, proto.TaskSysStat, statTaskTimeout)
		if err != nil {
			// Cached too, so pages queued behind this one don't each wait it out.
			res = proto.TaskResult{Error: err.Error()}
		}
		e.res, e.at = oldAgentError(res), time.Now()
	}
	writeJSON(w, e.res)
}

// POST /api/panel/agents/{id}/speed - have the agent measure its link to us.
func (s *Server) handleAgentSpeed(w http.ResponseWriter, r *http.Request) {
	id, ok := s.onlineAgent(w, r)
	if !ok {
		return
	}
	res, err := s.runTask(id, proto.TaskSysSpeed, speedTaskTimeout)
	if err != nil {
		writeJSON(w, proto.TaskResult{Error: err.Error()})
		return
	}
	writeJSON(w, oldAgentError(res))
}

// speedBlock is one chunk of incompressible bytes, made once and sent over and
// over, so a speed test costs the server no CPU beyond the copy.
var speedBlock = sync.OnceValue(func() []byte {
	b := make([]byte, 64<<10)
	rand.Read(b)
	return b
})

// GET /api/agent/speed/down?bytes=N - stream N bytes (default: as much as fits
// in the time window) for the agent to time. bytes=0 is the ping probe.
func (s *Server) handleSpeedDown(w http.ResponseWriter, r *http.Request) {
	if s.authenticateAgent(r) == nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	left := speedDownCap
	if v := r.URL.Query().Get("bytes"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n >= 0 && n < left {
			left = n
		}
	}
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("Cache-Control", "no-store")
	// Keep a proxy in front from buffering or compressing the stream.
	w.Header().Set("X-Accel-Buffering", "no")
	block := speedBlock()
	deadline := time.Now().Add(speedDownFor)
	for left > 0 && time.Now().Before(deadline) && r.Context().Err() == nil {
		chunk := block
		if len(chunk) > left {
			chunk = chunk[:left]
		}
		n, err := w.Write(chunk)
		if err != nil {
			return
		}
		left -= n
	}
}

// POST /api/agent/speed/up - swallow the agent's upload.
func (s *Server) handleSpeedUp(w http.ResponseWriter, r *http.Request) {
	if s.authenticateAgent(r) == nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	_, _ = io.Copy(io.Discard, http.MaxBytesReader(w, r.Body, speedUpCap))
	w.WriteHeader(http.StatusNoContent)
}