3x-ui-agent/internal/web/service/probenotify/sweep.go

package probenotify

import (
	"encoding/json"
	"fmt"
	"math/big"
	"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"
)

// nodeInfo is one agent as the sweep sees it: identity and freshness
// from the panel store, operations metadata and the traffic ledger
// from probe.db.
type nodeInfo struct {
	id          int64
	name        string
	lastSeenSec int64 // 0 = never reported
	meta        probe.Meta
	traffic     probe.Traffic
	hasTraffic  bool
}

// Seams the production wiring fills with the real stores; tests
// substitute fakes. Everything the sweep records goes through these.
var (
	loadNodesFn = loadNodes
	stateGetFn  = probe.StateGet
	stateSetFn  = probe.StateSet
	saveMetaFn  = probe.SaveMeta
)

func loadNodes() ([]nodeInfo, error) {
	db := database.GetDB()
	var agents []model.AgentNode
	if err := db.Find(&agents).Error; err != nil {
		return nil, err
	}
	nodeNames := map[int]string{}
	var nodes []model.Node
	if err := db.Find(&nodes).Error; err == nil {
		for _, n := range nodes {
			nodeNames[n.Id] = n.Name
		}
	}
	metas, _ := probe.MetaSnapshot()
	traffics, _ := probe.TrafficSnapshot()
	out := make([]nodeInfo, 0, len(agents))
	for _, a := range agents {
		name := a.Name
		if a.NodeID != nil {
			if bound, ok := nodeNames[*a.NodeID]; ok && bound != "" {
				name = bound
			}
		}
		meta, ok := metas[int64(a.Id)]
		if !ok {
			meta = probe.DefaultMeta(int64(a.Id))
		}
		traffic, hasTraffic := traffics[int64(a.Id)]
		out = append(out, nodeInfo{
			id:          int64(a.Id),
			name:        name,
			lastSeenSec: a.LastSeen / 1000,
			meta:        meta,
			traffic:     traffic,
			hasTraffic:  hasTraffic,
		})
	}
	return out, nil
}

func (n nodeInfo) online(nowSec int64) bool {
	return n.lastSeenSec > 0 && nowSec-n.lastSeenSec <= onlineWindow
}

// watch is the sweep state that need not outlive the process (hub
// Watch): when a sweep first found each node absent, and when each
// node last returned from a long absence. The durable half of the
// offline bookkeeping — the announced mark — is in probe.db.
type watch struct {
	mu       sync.Mutex
	absent   map[int64]int64
	returned map[int64]int64
}

func newWatch() *watch {
	return &watch{absent: map[int64]int64{}, returned: map[int64]int64{}}
}

var defaultWatch = newWatch()

func downKey(agentID int64) string { return "notify_down:" + strconv.FormatInt(agentID, 10) }

// Sweep runs one connectivity + traffic pass and queues what it
// raises. Called by the 30s job.
func Sweep() {
	for _, note := range runSweep(time.Now(), defaultWatch) {
		Send(note)
	}
}

// runSweep is one pass over every node (hub sweep()). Offline and
// returning nodes are each gathered into one alert per pass, so an
// outage on the master's side — which drops every agent within about
// 30s — produces one or two alerts rather than one per node. A return
// is announced only after an offline alert, whether or not the node
// still has alerts enabled. Nothing is recorded as announced while
// no channel is configured, so configuring one afterwards still
// reports nodes already down and traffic already past a step.
func runSweep(now time.Time, w *watch) []Note {
	nodes, err := loadNodesFn()
	if err != nil {
		logger.Warningf("probenotify: sweep load nodes: %v", err)
		return nil
	}
	nowSec := now.Unix()
	ids := make(map[int64]bool, len(nodes))
	for _, n := range nodes {
		ids[n.id] = true
	}

	w.mu.Lock()
	for id := range w.absent {
		if !ids[id] {
			delete(w.absent, id)
		}
	}
	for id := range w.returned {
		if !ids[id] {
			delete(w.returned, id)
		}
	}
	w.mu.Unlock()

	grace := number("notify_grace") * 60
	armed := len(channelsFn()) > 0

	// The announced offline marks. A failed read skips the whole
	// connectivity section: taken for "not announced", it would
	// re-raise alerts the state says already went out.
	marks := make(map[int64]int64, len(nodes))
	marksOK := true
	for _, n := range nodes {
		v, ok, err := stateGetFn(downKey(n.id))
		if err != nil {
			logger.Warningf("probenotify: read offline mark %d: %v", n.id, err)
			marksOK = false
			break
		}
		if ok {
			if ts, perr := strconv.ParseInt(v, 10, 64); perr == nil {
				marks[n.id] = ts
			}
		}
	}

	var notes []Note
	if marksOK {
		var down, up []batchItem
		aborted := false
		for _, n := range nodes {
			if n.online(nowSec) {
				w.mu.Lock()
				if since, was := w.absent[n.id]; was {
					delete(w.absent, n.id)
					if nowSec-since >= grace {
						w.returned[n.id] = nowSec
					}
				}
				w.mu.Unlock()
				if ds := marks[n.id]; ds > 0 {
					if err := stateSetFn(downKey(n.id), ""); err != nil {
						logger.Warningf("probenotify: clear offline mark %d: %v", n.id, err)
						aborted = true
						break
					}
					delete(marks, n.id)
					up = append(up, batchItem{n.name, "离线 " + span(nowSec-ds)})
				}
				continue
			}
			// A node that has never reported cannot be announced,
			// and tracking its absence would count its first
			// connection as a return from an outage.
			if n.lastSeenSec == 0 {
				continue
			}
			w.mu.Lock()
			since, seen := w.absent[n.id]
			if !seen {
				since = nowSec
				w.absent[n.id] = nowSec
			}
			flapping := false
			if back, ok := w.returned[n.id]; ok && since-back < flapWindow {
				flapping = true
			}
			w.mu.Unlock()
			wait := grace
			if flapping {
				wait = flapGrace
			}
			if armed && n.meta.Notify && marks[n.id] == 0 && nowSec-since >= wait {
				if err := stateSetFn(downKey(n.id), strconv.FormatInt(n.lastSeenSec, 10)); err != nil {
					logger.Warningf("probenotify: set offline mark %d: %v", n.id, err)
					aborted = true
					break
				}
				marks[n.id] = n.lastSeenSec
				down = append(down, batchItem{n.name, "最后上报 " + clock(n.lastSeenSec)})
			}
		}
		if !aborted {
			if note, ok := batch("offline", "🔴", "离线", down); ok {
				notes = append(notes, note)
			}
			if note, ok := batch("online", "🟢", "恢复在线", up); ok {
				notes = append(notes, note)
			}
		}
	}

	notes = append(notes, trafficNotes(now, nodes, ids, armed)...)
	return notes
}

// alertStep is one node's announced traffic step in its billing
// period, persisted as JSON under "traffic_alerted" so a step
// crossed while the master was stopped is announced on its return
// and one announced before is not repeated (hub TRAFFIC_ALERTED).
// The wire shape matches the hub's: {"<id>": ["<period>", <step>]}.
type alertStep struct {
	Period string
	Step   int64
}

func (s alertStep) MarshalJSON() ([]byte, error) {
	return json.Marshal([]any{s.Period, s.Step})
}

func (s *alertStep) UnmarshalJSON(data []byte) error {
	var pair []json.RawMessage
	if err := json.Unmarshal(data, &pair); err != nil {
		return err
	}
	if len(pair) != 2 {
		return fmt.Errorf("traffic step: want [period, step]")
	}
	if err := json.Unmarshal(pair[0], &s.Period); err != nil {
		return err
	}
	return json.Unmarshal(pair[1], &s.Step)
}

func reached(used, limit, percent int64) bool {
	// i128 in the hub: 100 × a limit near the int64 ceiling would
	// overflow, so compare in big integers (this runs per node per
	// sweep, dozens at most).
	u := new(big.Int).Mul(big.NewInt(used), big.NewInt(100))
	l := new(big.Int).Mul(big.NewInt(limit), big.NewInt(percent))
	return u.Cmp(l) >= 0
}

// trafficNotes is the sweep's traffic section. Read and written as
// one state value every pass, so a restored database brings its own
// steps. Unset/unparseable state means "not primed": the steps
// already reached are recorded rather than announced, so an upgrade
// does not re-announce every crossing (hub).
func trafficNotes(now time.Time, nodes []nodeInfo, ids map[int64]bool, armed bool) []Note {
	percent := number("notify_traffic")
	if percent <= 0 || !armed {
		return nil
	}
	stored, hasStored, err := stateGetFn("traffic_alerted")
	if err != nil {
		// A failed read skips the section: taken for "nothing
		// stored", it would record this sweep's crossings without
		// announcing them.
		logger.Warningf("probenotify: reading the traffic steps announced failed: %v", err)
		return nil
	}
	alerted := map[int64]alertStep{}
	primed := false
	if hasStored {
		var parsed map[int64]alertStep
		if jerr := json.Unmarshal([]byte(stored), &parsed); jerr == nil {
			alerted = parsed
			primed = true
		}
	}
	for id := range alerted {
		if !ids[id] {
			delete(alerted, id)
		}
	}
	var metered []batchItem
	for _, n := range nodes {
		if n.meta.TrafficLimit <= 0 || !n.hasTraffic {
			continue
		}
		used := n.traffic.MonthUsed(n.meta.TrafficMode)
		// The hub's ledger read reports the CURRENT period start as
		// month_start (its gate zeroes the counters of a stale
		// period); our gate keeps the stored stamp, so the current
		// period is computed here and the gated counters already
		// agree with it.
		period := probe.PeriodStart(dateOnly(now), n.meta.TrafficResetDay).Format("2006-01-02")
		step := int64(0)
		switch {
		case reached(used, n.meta.TrafficLimit, 100):
			step = 100
		case reached(used, n.meta.TrafficLimit, percent):
			step = percent
		}
		// A step below the one recorded is a manual correction;
		// storing it lets the same crossing be announced again.
		last, existed := alerted[n.id]
		alerted[n.id] = alertStep{Period: period, Step: step}
		told := existed && last.Period == period && last.Step >= step
		if step == 0 || told || !primed {
			continue
		}
		mode := switchModeName(n.meta.TrafficMode)
		share := new(big.Int).Mul(big.NewInt(used), big.NewInt(100))
		share.Div(share, big.NewInt(n.meta.TrafficLimit))
		metered = append(metered, batchItem{n.name, fmt.Sprintf(
			"已用 %s%%，%s / %s（%s），本期自 %s 起",
			share.String(), gib(used), gib(n.meta.TrafficLimit), mode, period)})
	}
	// Announced only once recorded: unrecorded, the next sweep would
	// announce the same step again, every 30 seconds while writes
	// fail.
	state, merr := json.Marshal(alerted)
	if merr != nil {
		logger.Warningf("probenotify: encoding the traffic steps failed: %v", merr)
		return nil
	}
	if !hasStored || stored != string(state) {
		if err := stateSetFn("traffic_alerted", string(state)); err != nil {
			logger.Warningf("probenotify: saving the traffic steps announced failed: %v", err)
			return nil
		}
	}
	if note, ok := batch("traffic", "⚠️", "流量提醒", metered); ok {
		return []Note{note}
	}
	return nil
}

func switchModeName(mode string) string {
	switch mode {
	case "up":
		return "仅上行"
	case "down":
		return "仅下行"
	case "max":
		return "取较大值"
	default:
		return "上下行相加"
	}
}

// Housekeeping runs the hourly half: roll online nodes' expiry dates
// forward first, then the daily expiry digest over the dates as they
// now stand (hub housekeeping order).
func Housekeeping() {
	for _, note := range runHousekeeping(time.Now()) {
		Send(note)
	}
}

func runHousekeeping(now time.Time) []Note {
	nodes, err := loadNodesFn()
	if err != nil {
		logger.Warningf("probenotify: housekeeping load nodes: %v", err)
		return nil
	}
	nowSec := now.Unix()
	today := dateOnly(now)
	var notes []Note

	// Auto-renewal: an online node past its expiry date is still
	// being paid for in practice — roll the date whole cycles into
	// the future and say so. Dates are the operator's local calendar
	// days (hub renew_online_nodes).
	var rolled []batchItem
	for _, n := range nodes {
		if !n.online(nowSec) || n.meta.ExpiresAt == nil {
			continue
		}
		expires, ok := parseDate(*n.meta.ExpiresAt)
		if !ok {
			continue
		}
		months, ok := cycleMonths(n.meta.BillingCycle)
		if !ok {
			continue
		}
		next := expires
		for next.Before(today) {
			next = addMonthsClamped(next, months)
		}
		if next.Equal(expires) {
			continue
		}
		meta := n.meta
		nextStr := next.Format("2006-01-02")
		meta.ExpiresAt = &nextStr
		if err := saveMetaFn(meta); err != nil {
			logger.Warningf("probenotify: roll expiry of %d: %v", n.id, err)
			continue
		}
		logger.Infof("probenotify: node %s is still up past %s, expiry rolled to %s", n.name, *n.meta.ExpiresAt, nextStr)
		rolled = append(rolled, batchItem{n.name, fmt.Sprintf("%s → %s", expires.Format("2006-01-02"), nextStr)})
	}
	if number("notify_expiry") > 0 {
		if note, ok := batch("renew", "🔁", "已自动续期", rolled); ok {
			notes = append(notes, note)
		}
	}

	if note, ok := expiryDigest(now, today, nodes); ok {
		notes = append(notes, note)
	}
	return notes
}

// expiryDigest is the once-a-day list of nodes expiring within the
// configured window, sent from 09:00 local. The date is stored once
// sent, so a restart does not repeat it — and not while no channel
// is configured, so one configured later that day still receives it
// (hub expiry_digest).
func expiryDigest(now, today time.Time, nodes []nodeInfo) (Note, bool) {
	days := number("notify_expiry")
	todayStr := today.Format("2006-01-02")
	if days == 0 || now.Hour() < 9 {
		return Note{}, false
	}
	if v, ok, err := stateGetFn("notify_expiry_sent"); err != nil {
		logger.Warningf("probenotify: read expiry digest date: %v", err)
		return Note{}, false
	} else if ok && v == todayStr {
		return Note{}, false
	}
	if len(channelsFn()) == 0 {
		return Note{}, false
	}
	type soonItem struct {
		date time.Time
		name string
	}
	var soon []soonItem
	for _, n := range nodes {
		if n.meta.ExpiresAt == nil {
			continue
		}
		date, ok := parseDate(*n.meta.ExpiresAt)
		if !ok {
			continue
		}
		diff := int(date.Sub(today).Hours() / 24)
		if diff >= 0 && diff <= int(days) {
			soon = append(soon, soonItem{date, n.name})
		}
	}
	sort.Slice(soon, func(i, j int) bool {
		if !soon[i].date.Equal(soon[j].date) {
			return soon[i].date.Before(soon[j].date)
		}
		return soon[i].name < soon[j].name
	})
	if err := stateSetFn("notify_expiry_sent", todayStr); err != nil {
		logger.Warningf("probenotify: save expiry digest date: %v", err)
		return Note{}, false
	}
	items := make([]batchItem, 0, len(soon))
	for _, s := range soon {
		diff := int(s.date.Sub(today).Hours() / 24)
		left := "今天到期"
		if diff > 0 {
			left = fmt.Sprintf("还剩 %d 天", diff)
		}
		items = append(items, batchItem{s.name, fmt.Sprintf("%s %s", s.date.Format("2006-01-02"), left)})
	}
	return batch("expiry", "⏳", "即将到期", items)
}

// cycleMonths resolves a stored billing cycle to its length in
// months (hub cycle_months): the named cycles, or "{n}m" with
// 1 ≤ n ≤ 1200. "once" and anything unknown do not renew.
func cycleMonths(cycle string) (int64, bool) {
	switch cycle {
	case "monthly":
		return 1, true
	case "quarterly":
		return 3, true
	case "semiannual":
		return 6, true
	case "yearly":
		return 12, true
	case "biennial":
		return 24, true
	case "triennial":
		return 36, true
	}
	if strings.HasSuffix(cycle, "m") {
		if n, err := strconv.ParseInt(strings.TrimSuffix(cycle, "m"), 10, 64); err == nil && n >= 1 && n <= 1200 {
			return n, true
		}
	}
	return 0, false
}

// addMonthsClamped adds whole months, clamping the day to the target
// month's last day (Rust checked_add_months semantics; Go's AddDate
// would overflow Jan 31 + 1 month into March).
func addMonthsClamped(t time.Time, months int64) time.Time {
	y, m, d := t.Date()
	total := int64(int(m)-1) + months
	ny, nm := y+int(total/12), time.Month(total%12+1)
	if last := daysIn(ny, nm); d > last {
		d = last
	}
	return time.Date(ny, nm, d, 0, 0, 0, 0, t.Location())
}

func daysIn(year int, month time.Month) int {
	return time.Date(year, month+1, 0, 0, 0, 0, 0, time.UTC).Day()
}

func dateOnly(t time.Time) time.Time {
	y, m, d := t.Date()
	return time.Date(y, m, d, 0, 0, 0, 0, t.Location())
}

func parseDate(s string) (time.Time, bool) {
	t, err := time.ParseInLocation("2006-01-02", s, time.Local)
	if err != nil {
		return time.Time{}, false
	}
	return t, true
}