awg-profiler-golang/awg_profiler_test.go

package main

import (
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"
)

func TestRandMagicRange(t *testing.T) {
	for i := 0; i < 100000; i++ {
		v := randMagic()
		if v < 5 || v > 2147483647 {
			t.Fatalf("randMagic out of range: %d", v)
		}
	}
}

func TestRandRange(t *testing.T) {
	for i := 0; i < 100000; i++ {
		v := randRange(8, 32)
		if v < 8 || v > 32 {
			t.Fatalf("randRange out of [8,32]: %d", v)
		}
	}
	// single-value range
	if v := randRange(15, 15); v != 15 {
		t.Fatalf("randRange(15,15)=%d, want 15", v)
	}
	// inverted range returns min (with a warning)
	if v := randRange(50, 10); v != 50 {
		t.Fatalf("randRange(50,10)=%d, want 50", v)
	}
}

func TestIncrementIP(t *testing.T) {
	cases := []struct{ cur, net, want string }{
		{"", "10.0.0.0/24", "10.0.0.2"},
		{"10.0.0.2", "10.0.0.0/24", "10.0.0.3"},
		{"192.168.5.9", "192.168.5.0/24", "192.168.5.10"},
	}
	for _, c := range cases {
		if got := incrementIP(c.cur, c.net); got != c.want {
			t.Errorf("incrementIP(%q,%q)=%q, want %q", c.cur, c.net, got, c.want)
		}
	}
}

func TestServerAndFirstIP(t *testing.T) {
	if got := serverIP("10.110.90.0/24"); got != "10.110.90.1" {
		t.Errorf("serverIP=%q", got)
	}
	if got := getFirstClientIP("10.110.90.0/24"); got != "10.110.90.2" {
		t.Errorf("getFirstClientIP=%q", got)
	}
}

func TestSanitize(t *testing.T) {
	cases := map[string]string{
		"phone":       "phone",
		"my client!":  "my_client_",
		"a/b\\c":      "a_b_c",
		"ok.name-1_2": "ok.name-1_2",
	}
	for in, want := range cases {
		if got := sanitize(in); got != want {
			t.Errorf("sanitize(%q)=%q, want %q", in, got, want)
		}
	}
}

func TestConfigRoundTrip(t *testing.T) {
	dir := t.TempDir()
	configFile = filepath.Join(dir, "awg_config")

	in := &Config{
		Network: "10.0.0.0/24", Interface: "awg0", Port: "51820",
		PublicIP: "203.0.113.9", ServerPriv: "PRIV==", ServerPub: "PUB==",
		DNS: "1.1.1.1", MTU: "1420",
		Jc: "6", Jmin: "20", Jmax: "120", S1: "90", S2: "100",
		H1: "11", H2: "22", H3: "33", H4: "44",
	}
	writeConfig(in)

	out := loadConfig()
	if *out != *in {
		t.Fatalf("round-trip mismatch:\n in=%+v\nout=%+v", *in, *out)
	}

	// The written file must remain shell-sourceable (KEY="value" form).
	data, _ := os.ReadFile(configFile)
	if want := `SERVER_INTERFACE="awg0"`; !contains(string(data), want) {
		t.Errorf("config missing %q", want)
	}
}

func TestHidePrivateKey(t *testing.T) {
	in := "[Interface]\nPrivateKey = SECRETKEY==\nAddress = 10.0.0.1/24\n"
	out := hidePrivateKey(in)
	if contains(out, "SECRETKEY") {
		t.Errorf("private key not hidden: %q", out)
	}
	if !contains(out, "<hidden>") {
		t.Errorf("expected <hidden> marker: %q", out)
	}
}

func TestStateRoundTrip(t *testing.T) {
	dir := t.TempDir()
	stateFile = filepath.Join(dir, "awg_state.json")

	if loadState().DepsInstalled {
		t.Fatal("fresh state should not report deps installed")
	}
	saveState(&State{DepsInstalled: true, OSID: "debian", OSVersion: "12"})
	got := loadState()
	if !got.DepsInstalled || got.OSID != "debian" {
		t.Fatalf("state round-trip failed: %+v", got)
	}
}

// TestStatsAccumulation exercises the durable, per-peer traffic accumulation:
// normal deltas, counter-reset detection, peers missing from a dump, survival of
// a restart (re-read from disk), and reset. Each sampleStats call round-trips
// through statsFile, so persistence is tested implicitly.
func TestStatsAccumulation(t *testing.T) {
	dir := t.TempDir()
	dataDir = dir
	statsFile = filepath.Join(dir, "awg_stats.json")

	const k = "PEERKEY="
	live := func(rx, tx int64) map[string]PeerStat {
		return map[string]PeerStat{k: {PublicKey: k, TransferRx: rx, TransferTx: tx}}
	}
	rec := func() StatRecord {
		statsLock.Lock()
		defer statsLock.Unlock()
		return loadStatsLocked().Peers[k]
	}

	// First sighting: counts from now, so totals start at 0 with the raw counter
	// captured as the baseline.
	sampleStats(live(100, 40))
	if r := rec(); r.TotalRx != 0 || r.TotalTx != 0 || r.LastRx != 100 || r.LastTx != 40 {
		t.Fatalf("first sighting: got %+v", r)
	}
	if rec().Since == 0 {
		t.Fatal("first sighting must set Since")
	}

	// Normal growth: +200 rx, +60 tx.
	sampleStats(live(300, 100))
	if r := rec(); r.TotalRx != 200 || r.TotalTx != 60 || r.LastRx != 300 {
		t.Fatalf("delta: got %+v", r)
	}

	// Counter reset (reading dropped): the whole current value is new traffic.
	sampleStats(live(40, 10))
	if r := rec(); r.TotalRx != 240 || r.TotalTx != 70 || r.LastRx != 40 {
		t.Fatalf("reset branch: got %+v", r)
	}

	// Peer missing from a non-empty dump must not touch the record.
	sampleStats(map[string]PeerStat{"OTHER=": {PublicKey: "OTHER="}})
	if r := rec(); r.TotalRx != 240 || r.LastRx != 40 {
		t.Fatalf("missing peer changed record: got %+v", r)
	}

	// Simulate an app restart where the interface counter restarted near zero:
	// the persisted total must keep growing, never drop.
	before := rec().TotalRx
	sampleStats(live(15, 5))
	if r := rec(); r.TotalRx != before+15 {
		t.Fatalf("post-restart total must not drop: before=%d got %+v", before, r)
	}

	// Reset rebaselines to the current live counter and zeroes the totals.
	resetStats(k, 15, 5)
	if r := rec(); r.TotalRx != 0 || r.TotalTx != 0 || r.LastRx != 15 || r.LastTx != 5 {
		t.Fatalf("reset: got %+v", r)
	}
	// A subsequent unchanged reading adds nothing.
	sampleStats(live(15, 5))
	if r := rec(); r.TotalRx != 0 || r.TotalTx != 0 {
		t.Fatalf("post-reset unchanged sample added traffic: got %+v", r)
	}

	// Delete removes the record.
	deleteStats(k)
	if _, ok := func() (StatRecord, bool) {
		statsLock.Lock()
		defer statsLock.Unlock()
		r, ok := loadStatsLocked().Peers[k]
		return r, ok
	}(); ok {
		t.Fatal("deleteStats left the record behind")
	}
}

// TestRequireSlash24 checks the CIDR guard added after review.md flagged that
// serverIP/incrementIP/awgConfHeader all hardcode /24 regardless of what the
// operator types in — so anything but a /24 must be rejected up front.
func TestRequireSlash24(t *testing.T) {
	for _, v := range []string{"10.0.0.0/24", "192.168.5.0/24"} {
		requireSlash24(v) // must not die
	}

	for _, v := range []string{
		"10.0.0.0/16",   // wrong prefix length
		"10.0.0.5/24",   // not the network base address
		"not-a-cidr",    // unparseable
		"2001:db8::/24", // IPv6, not IPv4
	} {
		func() {
			webMode = true // turns die() into a recoverable panic
			defer func() { webMode = false; recover() }()
			requireSlash24(v)
			t.Errorf("requireSlash24(%q) should have been rejected", v)
		}()
	}
}

// TestCsrfSafe checks the CSRF guard added after review.md flagged that
// Basic-auth-only endpoints were reachable via a blind cross-site <form>
// POST/DELETE. HTML forms can never set Content-Type: application/json, so
// requiring it on every state-changing request blocks exactly that attack.
func TestCsrfSafe(t *testing.T) {
	get := httptest.NewRequest(http.MethodGet, "/api/status", nil)
	if !csrfSafe(get) {
		t.Error("GET must always be csrf-safe")
	}

	postNoCT := httptest.NewRequest(http.MethodPost, "/api/server/stop", nil)
	if csrfSafe(postNoCT) {
		t.Error("POST without Content-Type must be rejected")
	}

	postForm := httptest.NewRequest(http.MethodPost, "/api/server/stop", nil)
	postForm.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	if csrfSafe(postForm) {
		t.Error("POST with a form content-type (the only kind an HTML form can send) must be rejected")
	}

	postJSON := httptest.NewRequest(http.MethodPost, "/api/server/stop", nil)
	postJSON.Header.Set("Content-Type", "application/json")
	if !csrfSafe(postJSON) {
		t.Error("POST with application/json must be accepted")
	}

	del := httptest.NewRequest(http.MethodDelete, "/api/clients/1", nil)
	del.Header.Set("Content-Type", "application/json")
	if !csrfSafe(del) {
		t.Error("DELETE with application/json must be accepted")
	}
}

func contains(s, sub string) bool {
	for i := 0; i+len(sub) <= len(s); i++ {
		if s[i:i+len(sub)] == sub {
			return true
		}
	}
	return false
}