awg-profiler-golang/stats.go

package main

import (
	"encoding/json"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

// PeerStat holds the live transfer/handshake counters for a single peer, parsed
// from `awg show <iface> dump`. Keys are the peer public keys.
type PeerStat struct {
	PublicKey       string `json:"public_key"`
	Endpoint        string `json:"endpoint"`
	LatestHandshake int64  `json:"latest_handshake"` // unix seconds, 0 = never
	TransferRx      int64  `json:"transfer_rx"`      // bytes received by server from peer
	TransferTx      int64  `json:"transfer_tx"`      // bytes sent by server to peer
	Keepalive       string `json:"keepalive"`
}

// onlineWindow is how recent a handshake must be for a peer to count as online.
// It must stay above AmneziaWG/WireGuard's rekey interval (~120s): a live but
// idle peer only refreshes its handshake on rekey, so a shorter window would
// falsely flip connected peers to offline. 150s is the tightest safe value —
// it makes a genuinely disconnected peer drop within ~30s of its last rekey
// instead of the previous 180s.
const onlineWindow = 150 * time.Second

// Online reports whether the peer handshaked within onlineWindow.
func (p PeerStat) Online() bool {
	if p.LatestHandshake == 0 {
		return false
	}
	return time.Since(time.Unix(p.LatestHandshake, 0)) <= onlineWindow
}

// peerStats runs `awg show <iface> dump` and returns a map keyed by peer public
// key. On any error (interface down, awg missing) it returns an empty map so the
// caller can still render the registry without live data.
//
// Dump line layout for a peer (tab-separated):
//
//	public-key  preshared-key  endpoint  allowed-ips  latest-handshake  rx  tx  keepalive
//
// The first line describes the interface itself and is skipped.
func peerStats(iface string) map[string]PeerStat {
	stats := map[string]PeerStat{}
	if iface == "" {
		return stats
	}
	out, err := output("awg", "show", iface, "dump")
	if err != nil {
		return stats
	}
	lines := strings.Split(strings.TrimSpace(out), "\n")
	for i, line := range lines {
		if i == 0 || strings.TrimSpace(line) == "" {
			continue // interface header / blank
		}
		f := strings.Fields(line)
		if len(f) < 8 {
			continue
		}
		hs, _ := strconv.ParseInt(f[4], 10, 64)
		rx, _ := strconv.ParseInt(f[5], 10, 64)
		tx, _ := strconv.ParseInt(f[6], 10, 64)
		endpoint := f[2]
		if endpoint == "(none)" {
			endpoint = ""
		}
		stats[f[0]] = PeerStat{
			PublicKey:       f[0],
			Endpoint:        endpoint,
			LatestHandshake: hs,
			TransferRx:      rx,
			TransferTx:      tx,
			Keepalive:       f[7],
		}
	}
	return stats
}

// ─── persistent, per-peer cumulative traffic ────────────────────────────────────
//
// The live PeerStat counters above are volatile: the WG interface resets them on
// every restart, and in the container the userspace amneziawg-go data plane
// restarts together with the app, so a plain read is zeroed after any restart.
// StatRecord keeps a durable running total per peer in awg_stats.json that only
// ever grows, by folding successive live readings in as deltas.

// StatRecord is the persisted, cumulative traffic tally for one peer.
type StatRecord struct {
	PublicKey string `json:"public_key"`
	Since     int64  `json:"since"`      // unix: when accumulation started / was last cleared
	UpdatedAt int64  `json:"updated_at"` // unix: last sample that moved the totals
	TotalRx   int64  `json:"total_rx"`   // accumulated bytes received from the peer
	TotalTx   int64  `json:"total_tx"`   // accumulated bytes sent to the peer
	LastRx    int64  `json:"last_rx"`    // last raw counter seen — the delta baseline
	LastTx    int64  `json:"last_tx"`
}

// StatsStore is the on-disk shape of awg_stats.json, keyed by peer public key.
type StatsStore struct {
	Peers map[string]StatRecord `json:"peers"`
}

// statsLock serialises every read-modify-write of the stats store. It is
// deliberately independent of opLock (which guards WG-mutating operations): the
// background sampler must never block behind a long install/init operation.
var statsLock sync.Mutex

// loadStatsLocked reads awg_stats.json. A missing/unreadable file yields an empty
// store; a corrupt file is preserved as awg_stats.json.bad and treated as empty,
// so a parse error can never silently wipe good data. Callers must hold statsLock.
func loadStatsLocked() *StatsStore {
	data, err := os.ReadFile(statsFile)
	if err != nil {
		return &StatsStore{Peers: map[string]StatRecord{}}
	}
	var s StatsStore
	if err := json.Unmarshal(data, &s); err != nil {
		warn("Stats file %s is invalid — backing it up as %s.bad and starting fresh", statsFile, statsFile)
		_ = os.Rename(statsFile, statsFile+".bad")
		return &StatsStore{Peers: map[string]StatRecord{}}
	}
	if s.Peers == nil {
		s.Peers = map[string]StatRecord{}
	}
	return &s
}

// saveStatsLocked writes the store atomically (temp file + rename). Stats are
// best-effort telemetry: on any failure it warns and returns rather than dying,
// so a full disk can never crash the server or a CLI operation. Holds statsLock.
func saveStatsLocked(s *StatsStore) {
	data, err := json.MarshalIndent(s, "", "  ")
	if err != nil {
		warn("Failed to encode stats: %v", err)
		return
	}
	if err := os.MkdirAll(dataDir, 0700); err != nil {
		warn("Failed to create data dir for stats: %v", err)
		return
	}
	tmp, err := os.CreateTemp(dataDir, "stats-*.tmp")
	if err != nil {
		warn("Failed to write stats: %v", err)
		return
	}
	tmpName := tmp.Name()
	if _, err := tmp.Write(data); err != nil {
		tmp.Close()
		os.Remove(tmpName)
		warn("Failed to write stats: %v", err)
		return
	}
	tmp.Close()
	if err := os.Rename(tmpName, statsFile); err != nil {
		os.Remove(tmpName)
		warn("Failed to write stats: %v", err)
	}
}

// accumulate folds a raw counter reading into a running total. A reading that is
// not below the baseline contributes its increase; a reading that dropped means
// the counter was reset (interface/app restart), so the whole current value is
// counted as new traffic. Returns the new total and the new baseline. Because
// every delta is non-negative, a total can only ever grow — a restart (which
// only lowers the live counter) can never reduce the persisted total.
func accumulate(total, last, cur int64) (int64, int64) {
	if cur >= last {
		return total + (cur - last), cur
	}
	return total + cur, cur
}

// snapshotLocked returns a copy of the store's records for lock-free rendering.
func snapshotLocked(s *StatsStore) map[string]StatRecord {
	out := make(map[string]StatRecord, len(s.Peers))
	for k, v := range s.Peers {
		out[k] = v
	}
	return out
}

// sampleStats folds one live reading into the persisted totals and returns a
// snapshot for the caller to render. Passing in an already-fetched live map lets
// callers avoid a second `awg show dump`. An empty live map (interface down /
// dump failed) is a no-op that still returns the current snapshot — crucially it
// never rebaselines or zeroes anything, and peers missing from a non-empty dump
// are left untouched too.
func sampleStats(live map[string]PeerStat) map[string]StatRecord {
	statsLock.Lock()
	defer statsLock.Unlock()
	s := loadStatsLocked()
	now := nowUnix()
	changed := false
	for pk, ps := range live {
		rec, ok := s.Peers[pk]
		if !ok {
			// First sighting: start counting from now, ignoring whatever the raw
			// counter already holds (that traffic predates tracking). "since" is
			// the mark the user sees.
			s.Peers[pk] = StatRecord{
				PublicKey: pk, Since: now, UpdatedAt: now,
				LastRx: ps.TransferRx, LastTx: ps.TransferTx,
			}
			changed = true
			continue
		}
		nrx, brx := accumulate(rec.TotalRx, rec.LastRx, ps.TransferRx)
		ntx, btx := accumulate(rec.TotalTx, rec.LastTx, ps.TransferTx)
		if nrx != rec.TotalRx || ntx != rec.TotalTx || brx != rec.LastRx || btx != rec.LastTx {
			rec.TotalRx, rec.LastRx = nrx, brx
			rec.TotalTx, rec.LastTx = ntx, btx
			rec.UpdatedAt = now
			s.Peers[pk] = rec
			changed = true
		}
	}
	if changed {
		saveStatsLocked(s)
	}
	return snapshotLocked(s)
}

// resetStats clears a peer's accumulated totals and rebaselines to the current
// live counter, so subsequent samples count only traffic from now on. "since" is
// set to now — the fresh mark from which stats accumulate again.
func resetStats(pubkey string, curRx, curTx int64) StatRecord {
	statsLock.Lock()
	defer statsLock.Unlock()
	s := loadStatsLocked()
	now := nowUnix()
	rec := StatRecord{PublicKey: pubkey, Since: now, UpdatedAt: now, LastRx: curRx, LastTx: curTx}
	s.Peers[pubkey] = rec
	saveStatsLocked(s)
	return rec
}

// deleteStats drops a peer's record (called when a client is deleted).
func deleteStats(pubkey string) {
	statsLock.Lock()
	defer statsLock.Unlock()
	s := loadStatsLocked()
	if _, ok := s.Peers[pubkey]; ok {
		delete(s.Peers, pubkey)
		saveStatsLocked(s)
	}
}