3x-ui-agent/internal/web/service/probe_public.go

package service

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	"net"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/mhsanaei/3x-ui/v3/internal/database"
	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
	"github.com/mhsanaei/3x-ui/v3/internal/logger"
	"github.com/mhsanaei/3x-ui/v3/internal/web/service/probe"
)

// Item 7 B3: the public status surface (monitor hub contract). The
// node views are assembled whitelist-style from three sources — the
// agent_nodes snapshot rows (facts + latest metrics, B1), the probe
// ledger/metadata (B2) and the bound panel node's name/remark — and
// rendered into per-audience snapshot frames shared by /api/nodes and
// /api/ws, exactly as the hub renders one frame per audience per tick.

// ProbeOnlineThresholdMs is how stale an agent's last heartbeat may be
// before the public page calls it offline. Heartbeats land every 60 s;
// 150 s absorbs one lost beat plus jitter.
const ProbeOnlineThresholdMs = 150_000

// ProbeSnapshotTTL mirrors the hub's 1900 ms frame reuse: just under
// the 2 s push interval, so a frame is built at most once per tick.
const ProbeSnapshotTTL = 1900 * time.Millisecond

// ---- public-surface settings (main settings KV, hub key names) ----

func probeSetting(key string) (string, bool) {
	var row model.Setting
	if err := database.GetDB().Where("key = ?", key).First(&row).Error; err != nil {
		return "", false
	}
	return row.Value, true
}

func setProbeSetting(key, value string) error {
	db := database.GetDB()
	var row model.Setting
	err := db.Where("key = ?", key).First(&row).Error
	if database.IsNotFound(err) {
		return db.Create(&model.Setting{Key: key, Value: value}).Error
	}
	if err != nil {
		return err
	}
	row.Value = value
	return db.Save(&row).Error
}

// ProbeFastIntervalSec is the agent fast-report cadence in seconds
// (V2.0.1 probe batch): runtime metrics ride a lightweight report at
// this interval instead of the 60 s heartbeat. Default 3, operator
// range 1–60 (setting key probe_fast_interval).
func ProbeFastIntervalSec() int {
	v, ok := probeSetting("probe_fast_interval")
	if !ok {
		return 3
	}
	n, err := strconv.Atoi(strings.TrimSpace(v))
	if err != nil || n < 1 {
		return 3
	}
	if n > 60 {
		return 60
	}
	return n
}

// ProbePublicPageEnabled reports the public-page master switch. Per
// the item-7 scope decision the page ships OFF until the operator
// turns it on (setting public_page = "on"); the hub's own default is
// on, the deviation is deliberate and recorded in the B3 report.
func ProbePublicPageEnabled() bool {
	v, _ := probeSetting("public_page")
	return v == "on"
}

// SetProbePublicPageEnabled flips the public-page master switch.
func SetProbePublicPageEnabled(on bool) error {
	if on {
		return setProbeSetting("public_page", "on")
	}
	return setProbeSetting("public_page", "off")
}

// ProbeSiteName is the status page's site name (hub default "Monitor").
func ProbeSiteName() string {
	if v, ok := probeSetting("site_name"); ok && v != "" {
		return v
	}
	return "Monitor"
}

// ProbeSiteURL is the hub's public URL setting ("site"), "" when unset.
func ProbeSiteURL() string {
	v, _ := probeSetting("site")
	return v
}

// ProbeGitHubConfigured mirrors /api/me's github flag: whether a
// github_client_id is configured. (GitHub login itself is not ported;
// the panel's own session is the only login.)
func ProbeGitHubConfigured() bool {
	v, _ := probeSetting("github_client_id")
	return strings.TrimSpace(v) != ""
}

// ProbeCurrentTheme is the selected theme short ("" = built-in default).
func ProbeCurrentTheme() string {
	v, _ := probeSetting("theme")
	return v
}

// ProbeThemeConfigValue reads a theme's saved config JSON verbatim
// (setting key theme_config:{short}), "{}" when never saved. A read
// failure is an error, not an empty config: a caller saving on top of
// what it read would otherwise erase every override (hub discipline).
func ProbeThemeConfigValue(short string) (string, error) {
	var row model.Setting
	err := database.GetDB().Where("key = ?", "theme_config:"+short).First(&row).Error
	if database.IsNotFound(err) {
		return "{}", nil
	}
	if err != nil {
		return "", err
	}
	if row.Value == "" {
		return "{}", nil
	}
	return row.Value, nil
}

// SaveProbeThemeConfigRaw stores a theme's config JSON verbatim.
func SaveProbeThemeConfigRaw(short, raw string) error {
	return setProbeSetting("theme_config:"+short, raw)
}

// ---- node view assembly ----

type probeNodeSource struct {
	agent      model.AgentNode
	name       string
	remark     string
	meta       probe.Meta
	traffic    probe.Traffic
	nowMs      int64
	now        time.Time
	online     bool
	lastSeen   int64 // unix seconds, 0 = never seen
	metricsOn  bool
	metricsObj map[string]any
}

func loadProbeNodeSources() []probeNodeSource {
	now := time.Now()
	nowMs := now.UnixMilli()

	var agents []model.AgentNode
	if err := database.GetDB().Order("id ASC").Find(&agents).Error; err != nil {
		logger.Warningf("probe public: load agents: %v", err)
		return nil
	}
	nodesByID := map[int]model.Node{}
	var nodes []model.Node
	if err := database.GetDB().Find(&nodes).Error; err == nil {
		for _, n := range nodes {
			nodesByID[n.Id] = n
		}
	}
	metas, err := probe.MetaSnapshot()
	if err != nil {
		logger.Warningf("probe public: load meta: %v", err)
		metas = map[int64]probe.Meta{}
	}
	traffic, err := probe.TrafficSnapshot()
	if err != nil {
		logger.Warningf("probe public: load traffic: %v", err)
		traffic = map[int64]probe.Traffic{}
	}

	out := make([]probeNodeSource, 0, len(agents))
	for i := range agents {
		a := agents[i]
		src := probeNodeSource{agent: a, name: a.Name, nowMs: nowMs, now: now}
		if a.NodeID != nil {
			if n, ok := nodesByID[*a.NodeID]; ok {
				if n.Name != "" {
					src.name = n.Name
				}
				src.remark = n.Remark
			}
		}
		if m, ok := metas[int64(a.Id)]; ok {
			src.meta = m
		} else {
			src.meta = probe.DefaultMeta(int64(a.Id))
		}
		src.traffic = traffic[int64(a.Id)]
		src.online = a.LastSeen > 0 && nowMs-a.LastSeen <= ProbeOnlineThresholdMs
		src.lastSeen = a.LastSeen / 1000
		if src.online && a.ProbeAt > 0 {
			src.metricsOn = true
			src.metricsObj = probeLiveMetrics(&a, src.traffic)
		}
		out = append(out, src)
	}
	sort.SliceStable(out, func(i, j int) bool {
		if out[i].meta.Sort != out[j].meta.Sort {
			return out[i].meta.Sort < out[j].meta.Sort
		}
		return out[i].agent.Id < out[j].agent.Id
	})
	return out
}

// probeLiveMetrics builds the whitelisted live metrics object
// (hub PUBLIC_METRICS) from the latest snapshot, with the four ledger
// figures injected (hub INJECTED — they never arrive from an agent).
func probeLiveMetrics(a *model.AgentNode, t probe.Traffic) map[string]any {
	return map[string]any{
		"uptime":     a.UptimeSecs,
		"cpu":        float64(a.ProbeCPU),
		"load":       []float64{float64(a.ProbeLoad1), float64(a.ProbeLoad5), float64(a.ProbeLoad15)},
		"mem_total":  a.ProbeMemTotal,
		"mem_used":   a.ProbeMemUsed,
		"swap_total": a.ProbeSwapTotal,
		"swap_used":  a.ProbeSwapUsed,
		"disk_total": a.ProbeDiskTotal,
		"disk_used":  a.ProbeDiskUsed,
		"net_rx":     a.ProbeNetRx,
		"net_tx":     a.ProbeNetTx,
		"tcp":        a.ProbeTCP,
		"udp":        a.ProbeUDP,
		"procs":      a.ProbeProcs,
		"total_rx":   t.TotalRx,
		"total_tx":   t.TotalTx,
		"month_rx":   t.MonthRx,
		"month_tx":   t.MonthTx,
	}
}

// expiresInDays counts whole calendar days from the local today to
// the expires date, as the hub does on its own calendar.
func expiresInDays(expires *string, now time.Time) any {
	if expires == nil {
		return nil
	}
	d, err := time.ParseInLocation("2006-01-02", *expires, time.Local)
	if err != nil {
		return nil
	}
	ny, nm, nd := now.Date()
	ey, em, ed := d.Date()
	days := int(time.Date(ey, em, ed, 0, 0, 0, 0, time.UTC).Unix()/86400 -
		time.Date(ny, nm, nd, 0, 0, 0, 0, time.UTC).Unix()/86400)
	return days
}

// buildProbeNodeView renders one node for one audience. The public
// shape is built by whitelist (hub node_view): address, hostname and
// the private remark exist only in the admin (full) rendering.
func buildProbeNodeView(src probeNodeSource, full bool) map[string]any {
	a := src.agent
	m := src.meta
	t := src.traffic

	var lastSeenAgo any
	if src.lastSeen > 0 {
		ago := src.now.Unix() - src.lastSeen
		if ago < 0 {
			ago = 0
		}
		lastSeenAgo = ago
	}
	var metrics any
	if src.metricsOn {
		obj := make(map[string]any, len(src.metricsObj)+1)
		for k, v := range src.metricsObj {
			obj[k] = v
		}
		if full {
			obj["iface"] = a.ProbeIface
		}
		metrics = obj
	}

	view := map[string]any{
		"id":                a.Id,
		"name":              src.name,
		"country":           m.CountryPin,
		"group":             m.GroupName,
		"public_remark":     m.PublicRemark,
		"sort":              m.Sort,
		"public":            m.Public,
		"online":            src.online,
		"last_seen":         src.lastSeen,
		"last_seen_ago":     lastSeenAgo,
		"metrics":           metrics,
		"os":                a.ProbeOS,
		"kernel":            a.ProbeKernel,
		"arch":              a.ProbeArch,
		"virt":              a.ProbeVirt,
		"cpu_name":          a.ProbeCPUName,
		"cpu_cores":         a.ProbeCPUCores,
		"mem_total":         a.ProbeMemTotal,
		"swap_total":        a.ProbeSwapTotal,
		"disk_total":        a.ProbeDiskTotal,
		"agent_version":     a.AgentVersion,
		"price":             m.Price,
		"currency":          m.Currency,
		"billing_cycle":     m.BillingCycle,
		"expires_at":        m.ExpiresAt,
		"expires_in":        expiresInDays(m.ExpiresAt, src.now),
		"traffic_limit":     m.TrafficLimit,
		"traffic_mode":      m.TrafficMode,
		"traffic_reset_day": m.TrafficResetDay,
		"total_rx":          t.TotalRx,
		"total_tx":          t.TotalTx,
		"month_rx":          t.MonthRx,
		"month_tx":          t.MonthTx,
		"month_used":        t.MonthUsed(m.TrafficMode),
		"month_start":       t.MonthStart,
		"day_rx":            t.DayRx,
		"day_tx":            t.DayTx,
	}
	if !full {
		return view
	}

	// Admin-only half: identity and addresses never leave the panel
	// in the anonymous rendering (they are simply never built there).
	held4, held6 := a.IPv4, a.IPv6
	exits := []string{a.SeenIPv4, a.SeenIPv6}
	shown := probeAddresses(exits, held4, held6, m.IPv4Pin, m.IPv6Pin)
	addrList := make([]map[string]string, 0, len(shown))
	for _, s := range shown {
		addrList = append(addrList, map[string]string{"address": s.addr, "source": s.source})
	}
	ip := ""
	for _, e := range exits {
		if e != "" {
			ip = e
			break
		}
	}
	var interval any
	if src.online {
		interval = 60 // heartbeat cadence is fixed at 60 s by design
	}
	view["hostname"] = a.ProbeHostname
	view["ip"] = ip
	view["ipv4"] = held4
	view["ipv6"] = held6
	view["ipv4_pin"] = m.IPv4Pin
	view["ipv6_pin"] = m.IPv6Pin
	view["addresses"] = addrList
	view["ipv4_auto"] = probeAutoAddress(exits, held4, held6, "", m.IPv6Pin, false)
	view["ipv6_auto"] = probeAutoAddress(exits, held4, held6, m.IPv4Pin, "", true)
	view["interval"] = interval
	view["country_pin"] = m.CountryPin
	view["country_auto"] = ""
	view["remark"] = src.remark
	view["notify"] = m.Notify
	return view
}

type probeShownAddr struct {
	addr   string
	source string
}

// probeAddresses is the hub's addresses() port: per family a manual
// pin wins, then a public interface address, then — when the interface
// holds only a private address of that family — the public exit the
// master observed. With nothing public anywhere it falls back to the
// held interface addresses, then to the connection source.
func probeAddresses(exits []string, held4, held6, pin4, pin6 string) []probeShownAddr {
	public := func(s string, v6 bool) bool {
		ip := net.ParseIP(s)
		if ip == nil || (ip.To4() == nil) != v6 {
			return false
		}
		return IsPublicIPLiteral(s)
	}
	family := func(pin, held string, v6 bool) *probeShownAddr {
		if pin != "" {
			return &probeShownAddr{pin, "manual"}
		}
		if public(held, v6) {
			return &probeShownAddr{held, "interface"}
		}
		if held == "" {
			return nil
		}
		for _, e := range exits {
			if public(e, v6) {
				return &probeShownAddr{e, "exit"}
			}
		}
		return nil
	}
	var shown []probeShownAddr
	if s := family(pin4, held4, false); s != nil {
		shown = append(shown, *s)
	}
	if s := family(pin6, held6, true); s != nil {
		shown = append(shown, *s)
	}
	if len(shown) > 0 {
		return shown
	}
	for _, h := range []string{held4, held6} {
		if h != "" {
			shown = append(shown, probeShownAddr{h, "interface"})
		}
	}
	if len(shown) > 0 {
		return shown
	}
	for _, e := range exits {
		if e != "" {
			return []probeShownAddr{{e, "connection"}}
		}
	}
	return nil
}

// probeAutoAddress recomputes one family's shown address with its own
// pin cleared (the edit form's fallback hint, hub ipv4_auto/ipv6_auto).
func probeAutoAddress(exits []string, held4, held6, pin4, pin6 string, v6 bool) string {
	for _, s := range probeAddresses(exits, held4, held6, pin4, pin6) {
		if strings.Contains(s.addr, ":") == v6 {
			return s.addr
		}
	}
	return ""
}

// BuildProbeNodeViews renders every node for one audience; anonymous
// renders only public=1 nodes.
func BuildProbeNodeViews(full bool) []map[string]any {
	sources := loadProbeNodeSources()
	views := make([]map[string]any, 0, len(sources))
	for _, src := range sources {
		if !full && !src.meta.Public {
			continue
		}
		views = append(views, buildProbeNodeView(src, full))
	}
	return views
}

// ---- snapshot frames (shared by /api/nodes and /api/ws) ----

type probeFrame struct {
	at   time.Time
	body []byte
	gz   []byte
}

var probeSnapshots struct {
	mu     sync.Mutex
	frames [2]probeFrame
}

// ProbeSnapshotJSON returns the audience's rendered snapshot frame,
// rebuilding it at most once per ProbeSnapshotTTL.
func ProbeSnapshotJSON(full bool) []byte {
	probeSnapshots.mu.Lock()
	defer probeSnapshots.mu.Unlock()
	idx := 0
	if full {
		idx = 1
	}
	f := &probeSnapshots.frames[idx]
	if f.body == nil || time.Since(f.at) < 0 || time.Since(f.at) >= ProbeSnapshotTTL {
		frame := struct {
			Nodes []map[string]any `json:"nodes"`
			Admin bool             `json:"admin"`
		}{Nodes: BuildProbeNodeViews(full), Admin: full}
		body, err := json.Marshal(frame)
		if err != nil {
			body = []byte(`{"nodes":[],"admin":false}`)
		}
		f.at = time.Now()
		f.body = body
		f.gz = nil
	}
	return f.body
}

// ProbeSnapshotGzip is the anonymous frame gzipped, built lazily once
// per frame (hub live_gzip) for /api/ws ?gzip binary frames and for
// pre-compressed /api/nodes answers.
func ProbeSnapshotGzip() []byte {
	body := ProbeSnapshotJSON(false)
	probeSnapshots.mu.Lock()
	defer probeSnapshots.mu.Unlock()
	f := &probeSnapshots.frames[0]
	if f.gz == nil {
		var buf bytes.Buffer
		w := gzip.NewWriter(&buf)
		_, _ = w.Write(body)
		_ = w.Close()
		f.gz = buf.Bytes()
	}
	return f.gz
}

// InvalidateProbeSnapshot drops both cached frames so the next read
// rebuilds (called after metadata edits, as the hub does on node
// writes).
func InvalidateProbeSnapshot() {
	probeSnapshots.mu.Lock()
	defer probeSnapshots.mu.Unlock()
	probeSnapshots.frames = [2]probeFrame{}
}

// ProbeNodeReadable mirrors the hub's readable(): the panel sees
// every node; an anonymous caller only public=1 nodes, and only while
// the public page is on. As on the hub, an authed caller may read a
// window for a node that does not exist (it comes back empty); the
// anonymous side needs the agent to exist with the public flag set,
// consistent with what /api/nodes lists.
func ProbeNodeReadable(full bool, agentID int64) bool {
	if full {
		return true
	}
	if !ProbePublicPageEnabled() || !ProbeAgentExists(agentID) {
		return false
	}
	meta, ok, err := probe.MetaFor(agentID)
	if err != nil {
		return false
	}
	if !ok {
		return probe.DefaultMeta(agentID).Public
	}
	return meta.Public
}

// ProbeAgentExists reports whether an agent row exists.
func ProbeAgentExists(agentID int64) bool {
	var count int64
	if err := database.GetDB().Model(&model.AgentNode{}).Where("id = ?", agentID).Count(&count).Error; err != nil {
		return false
	}
	return count > 0
}