3x-ui-agent/internal/web/service/probe/traffic.go

package probe

import (
	"database/sql"
	"math"
	"strings"
	"time"
)

// Traffic is one agent's whole-machine ledger row (hub Traffic): the
// running totals booked from the kernel counters of its counted
// interfaces. This is the machine-wide account (close to what providers
// bill), separate from the panel's per-inbound/per-client xray traffic.
type Traffic struct {
	TotalRx    int64  `json:"totalRx"`
	TotalTx    int64  `json:"totalTx"`
	MonthRx    int64  `json:"monthRx"`
	MonthTx    int64  `json:"monthTx"`
	MonthStart string `json:"monthStart"`
	DayRx      int64  `json:"dayRx"`
	DayTx      int64  `json:"dayTx"`
	DayStart   string `json:"-"` // read-side gating only, as in hub
}

// MonthUsed is this period's usage as the node's plan meters it
// (hub month_used): "up" counts tx only, "down" rx only, "max" the
// larger direction, anything else (incl. "sum") both directions.
func (t Traffic) MonthUsed(mode string) int64 {
	switch mode {
	case "up":
		return t.MonthTx
	case "down":
		return t.MonthRx
	case "max":
		if t.MonthRx > t.MonthTx {
			return t.MonthRx
		}
		return t.MonthTx
	default:
		return satAdd(t.MonthRx, t.MonthTx)
	}
}

func satAdd(a, b int64) int64 {
	if b > 0 && a > math.MaxInt64-b {
		return math.MaxInt64
	}
	if b < 0 && a < math.MinInt64-b {
		return math.MinInt64
	}
	return a + b
}

// PeriodStart is the start of the billing period containing the given
// date for a reset day of month (hub period_start). A reset day past
// the end of a short month lands on that month's last day.
func PeriodStart(today time.Time, resetDay int64) time.Time {
	if resetDay < 1 {
		resetDay = 1
	}
	if resetDay > 31 {
		resetDay = 31
	}
	clamped := func(y int, m time.Month) time.Time {
		last := time.Date(y, m+1, 1, 0, 0, 0, 0, time.Local).AddDate(0, 0, -1).Day()
		d := int(resetDay)
		if d > last {
			d = last
		}
		return time.Date(y, m, d, 0, 0, 0, 0, time.Local)
	}
	today = time.Date(today.Year(), today.Month(), today.Day(), 0, 0, 0, 0, time.Local)
	this := clamped(today.Year(), today.Month())
	if !today.Before(this) {
		return this
	}
	if today.Month() == time.January {
		return clamped(today.Year()-1, time.December)
	}
	return clamped(today.Year(), today.Month()-1)
}

const dateLayout = "2006-01-02"

func dateOnly(t time.Time) time.Time {
	return time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, time.Local)
}

// ensureRows lazily creates the agent's meta + traffic rows so booking
// and reading never special-case a first report. A meta row created
// here is also the moment the agent joins the probe system, so the
// auto-join probe tasks are assigned then — and only then, matching
// the hub, where auto_join applies to nodes created after the task
// was flagged, never retroactively to existing ones (B4).
func (s *Store) ensureRows(agentID int64) error {
	res, err := s.db.Exec(`INSERT OR IGNORE INTO probe_meta (agent_id) VALUES (?)`, agentID)
	if err != nil {
		return err
	}
	if n, err := res.RowsAffected(); err == nil && n > 0 {
		if _, err := s.db.Exec(`INSERT OR IGNORE INTO probe_ping_node (task_id, agent_id)
			SELECT id, ? FROM probe_ping_task WHERE auto_join`, agentID); err != nil {
			return err
		}
	}
	_, err = s.db.Exec(`INSERT OR IGNORE INTO probe_traffic (agent_id) VALUES (?)`, agentID)
	return err
}

// ResetDay reads the agent's billing reset day from probe_meta
// (default 1 when unset, hub's node default).
func (s *Store) ResetDay(agentID int64) (int64, error) {
	var day int64
	err := s.db.QueryRow(`SELECT traffic_reset_day FROM probe_meta WHERE agent_id = ?`, agentID).Scan(&day)
	if err == sql.ErrNoRows {
		return 1, nil
	}
	if err != nil {
		return 1, err
	}
	if day < 1 {
		return 1, nil
	}
	if day > 31 {
		return 31, nil
	}
	return day, nil
}

// accumulate books one reading of an agent's raw kernel counters into
// its running totals — a faithful port of hub db.rs accumulate().
//
// Only bytes the master observed a counter climb through are booked.
// Without a baseline under this exact boot there is nothing to subtract
// from: a first reading, a counter that shrank under the same boot, and
// a changed boot_id all book a zero delta and merely re-align the
// baseline (booking the bare reading would charge the machine's whole
// history again). The one exception is a machine that booted after the
// master last heard from it (booted > prevLastSeenSec): its counters
// started from zero while unobserved, so the reading is exactly the
// bytes since boot and is booked whole — a reboot during a master
// outage would otherwise lose the whole outage.
//
// at dates the reading for the day and billing period (hub local
// calendar); month/day boundaries never move a stored stamp backwards,
// so a panel correction stamped with the current period is not undone
// by a held-back reading.
func (s *Store) accumulate(agentID int64, bootID string, rx, tx int64, at time.Time, booted *int64, prevLastSeenSec int64) (Traffic, error) {
	if err := s.ensureRows(agentID); err != nil {
		return Traffic{}, err
	}
	resetDay, err := s.ResetDay(agentID)
	if err != nil {
		return Traffic{}, err
	}
	var (
		prevBoot, monthStart, dayStart                   string
		lastRx, lastTx                                   int64
		totalRx, totalTx, monthRx, monthTx, dayRx, dayTx int64
	)
	err = s.db.QueryRow(`
		SELECT boot_id, last_rx, last_tx, total_rx, total_tx, month_rx, month_tx,
		       month_start, day_rx, day_tx, day_start
		FROM probe_traffic WHERE agent_id = ?`, agentID).
		Scan(&prevBoot, &lastRx, &lastTx, &totalRx, &totalTx, &monthRx, &monthTx,
			&monthStart, &dayRx, &dayTx, &dayStart)
	if err != nil {
		return Traffic{}, err
	}

	kernelPart := func(b string) string {
		if i := strings.IndexByte(b, '/'); i >= 0 {
			return b[:i]
		}
		return b
	}
	unseenBoot := prevBoot != "" &&
		kernelPart(prevBoot) != kernelPart(bootID) &&
		booted != nil && prevLastSeenSec > 0 && *booted > prevLastSeenSec

	var dRx, dTx int64
	switch {
	case unseenBoot:
		dRx, dTx = rx, tx
	case prevBoot == "" || prevBoot != bootID:
		dRx, dTx = 0, 0
	default:
		if rx > lastRx {
			dRx = rx - lastRx
		}
		if tx > lastTx {
			dTx = tx - lastTx
		}
	}

	// Saturating adds: a total near MaxInt64 must not wrap negative.
	totalRx = satAdd(totalRx, dRx)
	totalTx = satAdd(totalTx, dTx)
	monthRx = satAdd(monthRx, dRx)
	monthTx = satAdd(monthTx, dTx)
	dayRx = satAdd(dayRx, dRx)
	dayTx = satAdd(dayTx, dTx)

	date := dateOnly(at)
	today := dateOnly(time.Now())
	dated := func(stored string) time.Time {
		if stamp, err := time.ParseInLocation(dateLayout, stored, time.Local); err == nil && !stamp.After(today) {
			if stamp.After(date) {
				return stamp
			}
			return date
		}
		return date
	}
	period := PeriodStart(dated(monthStart), resetDay)
	if monthStart != period.Format(dateLayout) {
		// A new billing period restarts the month counter, not the total.
		monthRx, monthTx = dRx, dTx
	}
	day := dated(dayStart)
	if dayStart != day.Format(dateLayout) {
		dayRx, dayTx = dRx, dTx
	}

	_, err = s.db.Exec(`
		UPDATE probe_traffic SET boot_id = ?, last_rx = ?, last_tx = ?,
			total_rx = ?, total_tx = ?, month_rx = ?, month_tx = ?, month_start = ?,
			day_rx = ?, day_tx = ?, day_start = ?
		WHERE agent_id = ?`,
		bootID, rx, tx, totalRx, totalTx, monthRx, monthTx, period.Format(dateLayout),
		dayRx, dayTx, day.Format(dateLayout), agentID)
	if err != nil {
		return Traffic{}, err
	}
	return Traffic{
		TotalRx: totalRx, TotalTx: totalTx,
		MonthRx: monthRx, MonthTx: monthTx, MonthStart: period.Format(dateLayout),
		DayRx: dayRx, DayTx: dayTx, DayStart: day.Format(dateLayout),
	}, nil
}

// SetTraffic corrects a ledger row from the panel side (hub
// set_traffic): month figures are re-stamped with the current period so
// the correction is not read back as belonging to an old period and
// discarded by the next booking. Nil fields keep their stored values.
// False when the agent has no meta row yet.
func (s *Store) SetTraffic(agentID int64, totalRx, totalTx, monthRx, monthTx *int64) (bool, error) {
	resetDay, err := s.ResetDay(agentID)
	if err != nil {
		return false, err
	}
	var exists int
	if err := s.db.QueryRow(`SELECT COUNT(*) FROM probe_meta WHERE agent_id = ?`, agentID).Scan(&exists); err != nil {
		return false, err
	}
	if exists == 0 {
		return false, nil
	}
	if err := s.ensureRows(agentID); err != nil {
		return false, err
	}
	period := PeriodStart(dateOnly(time.Now()), resetDay).Format(dateLayout)
	// COALESCE keeps stored values for nil fields; month figures only
	// move when the stored stamp already is the current period, and the
	// stamp always advances to the current period (hub semantics).
	_, err = s.db.Exec(`
		UPDATE probe_traffic SET
			total_rx = COALESCE(?, total_rx), total_tx = COALESCE(?, total_tx),
			month_rx = COALESCE(?, CASE WHEN month_start = ? THEN month_rx ELSE 0 END),
			month_tx = COALESCE(?, CASE WHEN month_start = ? THEN month_tx ELSE 0 END),
			month_start = ?
		WHERE agent_id = ?`,
		ptrOrNil(totalRx), ptrOrNil(totalTx), ptrOrNil(monthRx), period,
		ptrOrNil(monthTx), period, period, agentID)
	if err != nil {
		return false, err
	}
	return true, nil
}

func ptrOrNil(v *int64) any {
	if v == nil {
		return nil
	}
	return *v
}

// gate applies the read-side period gating of hub all_traffic: the
// period counters restart lazily in accumulate on the node's next
// report, so a node offline since before a boundary still holds the
// previous period's bytes on disk; readers see them as zero unless the
// stored stamp is the current period.
func gate(t Traffic, dayStart string, resetDay int64, now time.Time) Traffic {
	period := PeriodStart(dateOnly(now), resetDay).Format(dateLayout)
	if t.MonthStart != period {
		t.MonthRx, t.MonthTx = 0, 0
	}
	if dayStart != dateOnly(now).Format(dateLayout) {
		t.DayRx, t.DayTx = 0, 0
	}
	return t
}

// TrafficFor returns one agent's ledger with read-side period gating —
// the read port B3's /api/nodes injects total/month figures from.
func (s *Store) TrafficFor(agentID int64) (Traffic, bool, error) {
	var t Traffic
	var dayStart string
	err := s.db.QueryRow(`
		SELECT total_rx, total_tx, month_rx, month_tx, month_start, day_rx, day_tx, day_start
		FROM probe_traffic WHERE agent_id = ?`, agentID).
		Scan(&t.TotalRx, &t.TotalTx, &t.MonthRx, &t.MonthTx, &t.MonthStart, &t.DayRx, &t.DayTx, &dayStart)
	if err == sql.ErrNoRows {
		return Traffic{}, false, nil
	}
	if err != nil {
		return Traffic{}, false, err
	}
	resetDay, err := s.ResetDay(agentID)
	if err != nil {
		return Traffic{}, false, err
	}
	return gate(t, dayStart, resetDay, time.Now()), true, nil
}

// AllTraffic returns every agent's ledger in one query, gated as in
// TrafficFor — the bulk read port for B3's node list.
func (s *Store) AllTraffic() (map[int64]Traffic, error) {
	rows, err := s.db.Query(`
		SELECT t.agent_id, t.total_rx, t.total_tx, t.month_rx, t.month_tx, t.month_start,
		       t.day_rx, t.day_tx, t.day_start, COALESCE(m.traffic_reset_day, 1)
		FROM probe_traffic t LEFT JOIN probe_meta m ON m.agent_id = t.agent_id`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := make(map[int64]Traffic)
	now := time.Now()
	for rows.Next() {
		var id, resetDay int64
		var t Traffic
		var dayStart string
		if err := rows.Scan(&id, &t.TotalRx, &t.TotalTx, &t.MonthRx, &t.MonthTx, &t.MonthStart,
			&t.DayRx, &t.DayTx, &dayStart, &resetDay); err != nil {
			return nil, err
		}
		out[id] = gate(t, dayStart, resetDay, now)
	}
	return out, rows.Err()
}

// parseDay parses a stored "2006-01-02" stamp (used by tests).
func parseDay(s string) (time.Time, bool) {
	t, err := time.ParseInLocation(dateLayout, s, time.Local)
	if err != nil {
		return time.Time{}, false
	}
	return t, true
}