pangolin/server/internal/agentd/singbox_test.go

package agentd

import (
	"context"
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1"
)

// fakeRestarter counts Restart calls.
type fakeRestarter struct {
	mu sync.Mutex
	n  int
}

func (f *fakeRestarter) Restart(context.Context) error {
	f.mu.Lock()
	f.n++
	f.mu.Unlock()
	return nil
}

func (f *fakeRestarter) count() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.n
}

func testConfig(t *testing.T) Config {
	t.Helper()
	dir := t.TempDir()
	return Config{
		ControlPlaneAddr:  "bufnet",
		StateDir:          dir,
		SingboxConfigPath: filepath.Join(dir, "singbox", "config.json"),
		DeriveKey:         "test-derive-key",
		DebounceWindow:    20 * time.Millisecond,
		Insecure:          true,
	}.withDefaults()
}

func sampleSnapshot(creds ...*agentv1.Credential) *agentv1.ConfigSnapshot {
	return &agentv1.ConfigSnapshot{
		ConfigVersion: 7,
		Credentials:   creds,
		Reality: &agentv1.RealityInbound{
			ListenPort: 11443, PrivateKey: "pk", ShortID: "deadbeef",
			ServerName: "www.apple.com", HandshakeServer: "www.apple.com", HandshakePort: 443,
		},
		Hy2:           &agentv1.Hy2Inbound{ListenPort: 443, Masquerade: "https://www.bing.com", CertPath: "/c", KeyPath: "/k"},
		LastCommandID: 3,
	}
}

func TestRenderContainsCredentials(t *testing.T) {
	sb := NewSingBox(testConfig(t), nil)
	sb.ApplyConfig(sampleSnapshot(
		&agentv1.Credential{DpUUID: "aaaa", Protocol: agentv1.ProtocolBoth},
		&agentv1.Credential{DpUUID: "bbbb", Protocol: agentv1.ProtocolReality},
		&agentv1.Credential{DpUUID: "cccc", Protocol: agentv1.ProtocolHy2},
	), true)

	data, err := sb.RenderConfig()
	if err != nil {
		t.Fatal(err)
	}
	s := string(data)

	// REALITY inbound: aaaa + bbbb as VLESS uuids, with flow; cccc must NOT be there.
	for _, want := range []string{`"uuid": "aaaa"`, `"uuid": "bbbb"`, `"flow": "xtls-rprx-vision"`} {
		if !strings.Contains(s, want) {
			t.Errorf("rendered config missing %q\n%s", want, s)
		}
	}
	if strings.Contains(s, `"uuid": "cccc"`) {
		t.Errorf("hy2-only credential cccc leaked into REALITY users")
	}

	// Hy2 inbound: derived passwords for aaaa + cccc, never the raw dp_uuid as pw.
	var cfg map[string]any
	if err := json.Unmarshal(data, &cfg); err != nil {
		t.Fatalf("rendered config is not valid JSON: %v", err)
	}
	if !strings.Contains(s, DeriveHy2Password("aaaa", "test-derive-key")) {
		t.Error("hy2 password for aaaa not found")
	}
	if !strings.Contains(s, DeriveHy2Password("cccc", "test-derive-key")) {
		t.Error("hy2 password for cccc not found")
	}
}

func TestRevokeRemovesFromRender(t *testing.T) {
	sb := NewSingBox(testConfig(t), nil)
	sb.ApplyConfig(sampleSnapshot(&agentv1.Credential{DpUUID: "aaaa", Protocol: agentv1.ProtocolBoth}), true)
	sb.Revoke("aaaa")
	data, _ := sb.RenderConfig()
	if strings.Contains(string(data), "aaaa") {
		t.Errorf("revoked credential still present:\n%s", data)
	}
	if sb.Has("aaaa") {
		t.Error("Has reports revoked credential as present")
	}
}

func TestRotateGraceBothPresentThenOldExpires(t *testing.T) {
	cfg := testConfig(t)
	sb := NewSingBox(cfg, nil)
	now := time.Unix(1000, 0)
	sb.clock = func() time.Time { return now }
	sb.ApplyConfig(sampleSnapshot(&agentv1.Credential{DpUUID: "old", Protocol: agentv1.ProtocolBoth}), true)

	grace := int64(1005)
	sb.Rotate("old", &Cred{DpUUID: "new", Protocol: agentv1.ProtocolBoth}, grace)

	// During grace both coexist.
	if !sb.Has("old") || !sb.Has("new") {
		t.Fatalf("expected both old+new during grace; old=%v new=%v", sb.Has("old"), sb.Has("new"))
	}

	// After grace, sweep removes old, keeps new.
	now = time.Unix(1006, 0)
	removed := sb.sweepExpired()
	if len(removed) != 1 || removed[0] != "old" {
		t.Fatalf("expected old removed by sweep, got %v", removed)
	}
	if sb.Has("old") || !sb.Has("new") {
		t.Fatalf("after grace: old should be gone, new present; old=%v new=%v", sb.Has("old"), sb.Has("new"))
	}
}

func TestStateRecovery(t *testing.T) {
	cfg := testConfig(t)
	now := time.Unix(2000, 0)

	// First manager: add creds (one already expired), persist via mutation.
	sb1 := NewSingBox(cfg, nil)
	sb1.clock = func() time.Time { return now }
	sb1.Upsert(&Cred{DpUUID: "live", Protocol: agentv1.ProtocolBoth, ExpiresAt: 3000})
	sb1.Upsert(&Cred{DpUUID: "dead", Protocol: agentv1.ProtocolBoth, ExpiresAt: 1000}) // already expired

	// state.json must contain ONLY dp_uuid + expires_at — no identity fields.
	raw, err := os.ReadFile(cfg.StatePath())
	if err != nil {
		t.Fatal(err)
	}
	assertNoIdentityFields(t, raw)

	// Second manager (simulated kill -9 + restart) restores from disk.
	sb2 := NewSingBox(cfg, nil)
	sb2.clock = func() time.Time { return now }
	if err := sb2.LoadState(); err != nil {
		t.Fatal(err)
	}
	if !sb2.Has("live") {
		t.Error("live credential not recovered from state.json")
	}
	if sb2.Has("dead") {
		t.Error("expired credential should be dropped on load")
	}
	// Recovered creds default to BOTH + default flow until Register reconciles.
	for _, c := range sb2.Snapshot() {
		if c.DpUUID == "live" && (c.Protocol != agentv1.ProtocolBoth || c.Flow != DefaultFlow) {
			t.Errorf("recovered cred has unexpected defaults: %+v", c)
		}
	}
}

func TestTTLSweepReportsExpiry(t *testing.T) {
	cfg := testConfig(t)
	sb := NewSingBox(cfg, nil)
	now := time.Unix(5000, 0)
	sb.clock = func() time.Time { return now }

	var mu sync.Mutex
	var reported []string
	sb.OnExpire = func(dp string) { mu.Lock(); reported = append(reported, dp); mu.Unlock() }

	sb.Upsert(&Cred{DpUUID: "x", Protocol: agentv1.ProtocolBoth, ExpiresAt: 4999})
	sb.Upsert(&Cred{DpUUID: "y", Protocol: agentv1.ProtocolBoth, ExpiresAt: 9999})

	removed := sb.sweepExpired()
	if len(removed) != 1 || removed[0] != "x" {
		t.Fatalf("sweep removed = %v, want [x]", removed)
	}
	mu.Lock()
	defer mu.Unlock()
	if len(reported) != 1 || reported[0] != "x" {
		t.Fatalf("OnExpire reported = %v, want [x]", reported)
	}
}

func TestDebounceCoalescesRestarts(t *testing.T) {
	cfg := testConfig(t)
	cfg.DebounceWindow = 40 * time.Millisecond
	fr := &fakeRestarter{}
	sb := NewSingBox(cfg, fr)
	sb.ApplyConfig(sampleSnapshot(), true) // set inbounds so render succeeds

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go sb.Run(ctx)

	// Burst of changes within one debounce window.
	for i := 0; i < 5; i++ {
		sb.Upsert(&Cred{DpUUID: string(rune('a' + i)), Protocol: agentv1.ProtocolBoth})
	}
	time.Sleep(150 * time.Millisecond)

	if n := fr.count(); n != 1 {
		t.Fatalf("expected 1 coalesced restart, got %d", n)
	}
}

// assertNoIdentityFields fails if the blob contains any user-identity token.
func assertNoIdentityFields(t *testing.T, blob []byte) {
	t.Helper()
	lower := strings.ToLower(string(blob))
	for _, banned := range []string{"email", "user_id", "userid", "phone", "device_id", "deviceid", "account", "destination", "dns"} {
		if strings.Contains(lower, banned) {
			t.Errorf("forbidden identity field %q found in persisted/rendered data:\n%s", banned, blob)
		}
	}
}