alpine-router/traffic/tracker.go

package traffic

import (
	"bufio"
	"fmt"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	pollInterval     = 20 * time.Second
	OnlineWindow     = 5 * time.Minute
	trackerTableName = "nano-router-traffic"
)

type IPStats struct {
	TxBytes    uint64
	RxBytes    uint64
	LastActive time.Time
}

var (
	mu     sync.RWMutex
	stats  = map[string]*IPStats{}
	prev   = map[string][2]uint64{}
	useNFT bool
)

func Available() bool {
	mu.RLock()
	v := useNFT
	mu.RUnlock()
	return v
}

func Start() {
	if _, err := exec.LookPath("nft"); err != nil {
		return
	}
	if err := setupNFTTable(); err != nil {
		return
	}
	mu.Lock()
	useNFT = true
	mu.Unlock()

	go func() {
		poll()
		t := time.NewTicker(pollInterval)
		defer t.Stop()
		for range t.C {
			poll()
		}
	}()
}

func Get(ip string) IPStats {
	mu.RLock()
	defer mu.RUnlock()
	if s, ok := stats[ip]; ok {
		return *s
	}
	return IPStats{}
}

func IsOnline(ip string) bool {
	s := Get(ip)
	return !s.LastActive.IsZero() && time.Since(s.LastActive) < OnlineWindow
}

func EnsureIPTracked(ip string) {
	if ip == "" {
		return
	}
	mu.Lock()
	defer mu.Unlock()
	if _, ok := prev[ip]; ok {
		return
	}
	prev[ip] = [2]uint64{}
	stats[ip] = &IPStats{}
	addNFTRule(ip)
}

func setupNFTTable() error {
	exec.Command("nft", "delete", "table", "ip", trackerTableName).Run()

	script := fmt.Sprintf(`table ip %s {
	chain tx {
		type filter hook prerouting priority raw; policy accept;
	}
	chain rx {
		type filter hook postrouting priority raw; policy accept;
	}
}`, trackerTableName)

	cmd := exec.Command("nft", "-f", "-")
	cmd.Stdin = strings.NewReader(script)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("nft setup: %s: %w", strings.TrimSpace(string(out)), err)
	}
	return nil
}

func addNFTRule(ip string) {
	exec.Command("nft", "add", "rule", "ip", trackerTableName, "tx",
		"ip", "saddr", ip, "counter").Run()
	exec.Command("nft", "add", "rule", "ip", trackerTableName, "rx",
		"ip", "daddr", ip, "counter").Run()
}

var txCounterRe = regexp.MustCompile(`ip saddr (\S+) counter packets \d+ bytes (\d+)`)
var rxCounterRe = regexp.MustCompile(`ip daddr (\S+) counter packets \d+ bytes (\d+)`)

func poll() {
	mu.RLock()
	nft := useNFT
	mu.RUnlock()
	if !nft {
		return
	}

	current, err := readNFTCounters()
	if err != nil {
		return
	}

	now := time.Now()
	mu.Lock()
	defer mu.Unlock()

	for ip, cur := range current {
		p := prev[ip]
		dTx := delta(p[0], cur[0])
		dRx := delta(p[1], cur[1])

		s := stats[ip]
		if s == nil {
			s = &IPStats{}
			stats[ip] = s
		}
		s.TxBytes += dTx
		s.RxBytes += dRx
		if dTx > 0 || dRx > 0 {
			s.LastActive = now
		}
		prev[ip] = cur
	}

	for ip := range prev {
		if _, exists := current[ip]; !exists {
			prev[ip] = [2]uint64{0, 0}
			addNFTRule(ip)
		}
	}
}

func readNFTCounters() (map[string][2]uint64, error) {
	cmd := exec.Command("nft", "list", "table", "ip", trackerTableName)
	out, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	result := map[string][2]uint64{}
	scanner := bufio.NewScanner(strings.NewReader(string(out)))
	for scanner.Scan() {
		line := scanner.Text()
		if m := txCounterRe.FindStringSubmatch(line); m != nil {
			ip := m[1]
			b, _ := strconv.ParseUint(m[2], 10, 64)
			cur := result[ip]
			cur[0] += b
			result[ip] = cur
		} else if m := rxCounterRe.FindStringSubmatch(line); m != nil {
			ip := m[1]
			b, _ := strconv.ParseUint(m[2], 10, 64)
			cur := result[ip]
			cur[1] += b
			result[ip] = cur
		}
	}
	return result, nil
}

func delta(old, cur uint64) uint64 {
	if cur >= old {
		return cur - old
	}
	return cur
}