pangolin/server/internal/agentd/integration_test.go

package agentd

import (
	"context"
	"net"
	"os"
	"sort"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/wangjia/pangolin/server/internal/mtls"
	agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/status"
	"google.golang.org/grpc/test/bufconn"
)

const testNodeUUID = "node-test-uuid"

// mockCP is an in-memory control plane implementing AgentServiceServer.
type mockCP struct {
	ca *mtls.CA

	mu          sync.Mutex
	snapCreds   map[string]*agentv1.Credential // converged credential set for Register
	configVer   int64
	lastCmdID   int64
	issued      []*agentv1.Command // backlog for resume-after-drop
	acked       map[int64]bool
	registerN   int
	heartbeatN  int
	subReqs     []int64
	usage       []*agentv1.UsageReport
	resyncOnce  bool
	enrollToken string

	liveCh chan *agentv1.Command
	dropCh chan struct{}
}

func newMockCP(t *testing.T) *mockCP {
	t.Helper()
	dir := t.TempDir()
	ca, err := mtls.NewCA(mtls.CAConfig{
		KeyPath:  dir + "/ca.key",
		CertPath: dir + "/ca.crt",
	})
	if err != nil {
		t.Fatalf("new CA: %v", err)
	}
	return &mockCP{
		ca:          ca,
		snapCreds:   map[string]*agentv1.Credential{},
		acked:       map[int64]bool{},
		enrollToken: "test-token",
		liveCh:      make(chan *agentv1.Command),
		dropCh:      make(chan struct{}, 1),
	}
}

func (m *mockCP) Enroll(_ context.Context, req *agentv1.EnrollRequest) (*agentv1.EnrollResponse, error) {
	if req.BootstrapToken != m.enrollToken {
		return nil, status.Error(codes.Unauthenticated, "bad bootstrap token")
	}
	certPEM, err := m.ca.SignCSR(req.CSRPEM, testNodeUUID)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "sign csr: %v", err)
	}
	return &agentv1.EnrollResponse{
		NodeUUID:     testNodeUUID,
		CertPEM:      certPEM,
		CAPEM:        m.ca.CAPEM(),
		NotAfterUnix: time.Now().Add(90 * 24 * time.Hour).Unix(),
	}, nil
}

func (m *mockCP) Register(_ context.Context, _ *agentv1.RegisterRequest) (*agentv1.ConfigSnapshot, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.registerN++
	creds := make([]*agentv1.Credential, 0, len(m.snapCreds))
	for _, c := range m.snapCreds {
		creds = append(creds, c)
	}
	sort.Slice(creds, func(i, j int) bool { return creds[i].DpUUID < creds[j].DpUUID })
	return &agentv1.ConfigSnapshot{
		ConfigVersion: m.configVer,
		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: m.lastCmdID,
	}, nil
}

func (m *mockCP) Heartbeat(_ context.Context, _ *agentv1.HeartbeatRequest) (*agentv1.HeartbeatResponse, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.heartbeatN++
	resync := false
	if m.resyncOnce {
		m.resyncOnce = false
		resync = true
	}
	return &agentv1.HeartbeatResponse{NeedFullResync: resync, ServerTimeUnix: time.Now().Unix()}, nil
}

func (m *mockCP) Ack(_ context.Context, req *agentv1.AckRequest) (*agentv1.AckResponse, error) {
	m.mu.Lock()
	m.acked[req.CommandID] = true
	m.mu.Unlock()
	return &agentv1.AckResponse{}, nil
}

func (m *mockCP) ReportUsage(_ context.Context, req *agentv1.UsageReport) (*agentv1.UsageAck, error) {
	m.mu.Lock()
	m.usage = append(m.usage, req)
	m.mu.Unlock()
	return &agentv1.UsageAck{}, nil
}

func (m *mockCP) Subscribe(req *agentv1.SubscribeRequest, stream agentv1.AgentService_SubscribeServer) error {
	m.mu.Lock()
	m.subReqs = append(m.subReqs, req.LastCommandID)
	backlog := make([]*agentv1.Command, 0)
	for _, c := range m.issued {
		if c.CommandID > req.LastCommandID {
			backlog = append(backlog, c)
		}
	}
	m.mu.Unlock()

	for _, c := range backlog {
		if err := stream.Send(c); err != nil {
			return err
		}
	}
	for {
		select {
		case <-stream.Context().Done():
			return stream.Context().Err()
		case <-m.dropCh:
			return status.Error(codes.Unavailable, "simulated stream drop")
		case c := <-m.liveCh:
			if err := stream.Send(c); err != nil {
				return err
			}
		}
	}
}

// ─── mock helpers (test goroutine) ───────────────────────────────────────────

// seedCred adds a credential to the converged Register snapshot before start.
func (m *mockCP) seedCred(c *agentv1.Credential) {
	m.mu.Lock()
	m.snapCreds[c.DpUUID] = c
	if m.configVer == 0 {
		m.configVer = 1
	}
	m.mu.Unlock()
}

// push issues a live command AND folds it into the Register snapshot, so a later
// reconnect's Register stays consistent with the command stream.
func (m *mockCP) push(cmd *agentv1.Command) {
	m.applyToSnapshot(cmd, true) // updates snapCreds + lastCmdID
	m.mu.Lock()
	m.issued = append(m.issued, cmd)
	m.mu.Unlock()
	m.liveCh <- cmd
}

// enqueueOffline appends a command to the resume backlog WITHOUT touching the
// Register snapshot — so the only way the agent learns it is via Subscribe resume.
func (m *mockCP) enqueueOffline(cmd *agentv1.Command) {
	m.mu.Lock()
	// Intentionally does NOT advance lastCmdID: this command is not yet reflected
	// in the Register snapshot, so the agent must learn it via Subscribe resume.
	m.issued = append(m.issued, cmd)
	m.mu.Unlock()
}

func (m *mockCP) applyToSnapshot(cmd *agentv1.Command, bumpVer bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	switch cmd.Type {
	case agentv1.CommandTypeUpsert:
		m.snapCreds[cmd.Upsert.Credential.DpUUID] = cmd.Upsert.Credential
	case agentv1.CommandTypeRevoke:
		delete(m.snapCreds, cmd.Revoke.DpUUID)
	}
	if cmd.CommandID > m.lastCmdID {
		m.lastCmdID = cmd.CommandID
	}
	if bumpVer {
		m.configVer++
	}
}

func (m *mockCP) drop()                 { m.dropCh <- struct{}{} }
func (m *mockCP) registerCount() int    { m.mu.Lock(); defer m.mu.Unlock(); return m.registerN }
func (m *mockCP) heartbeatCount() int   { m.mu.Lock(); defer m.mu.Unlock(); return m.heartbeatN }
func (m *mockCP) isAcked(id int64) bool { m.mu.Lock(); defer m.mu.Unlock(); return m.acked[id] }
func (m *mockCP) subReqList() []int64 {
	m.mu.Lock()
	defer m.mu.Unlock()
	return append([]int64(nil), m.subReqs...)
}

// ─── test harness ─────────────────────────────────────────────────────────────

func startMock(t *testing.T, m *mockCP) func(ctx context.Context) (*grpc.ClientConn, error) {
	t.Helper()
	lis := bufconn.Listen(1024 * 1024)
	srv := grpc.NewServer()
	agentv1.RegisterAgentServiceServer(srv, m)
	go func() { _ = srv.Serve(lis) }()
	t.Cleanup(func() { srv.Stop() })
	return func(ctx context.Context) (*grpc.ClientConn, error) {
		return grpc.NewClient("passthrough:///bufnet",
			grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) {
				return lis.DialContext(ctx)
			}),
			grpc.WithTransportCredentials(insecure.NewCredentials()),
		)
	}
}

func fastConfig(t *testing.T) Config {
	cfg := testConfig(t)
	cfg.BootstrapToken = "test-token"
	cfg.HeartbeatInterval = 25 * time.Millisecond
	cfg.UsageInterval = time.Hour
	cfg.TTLScanInterval = time.Hour
	cfg.DebounceWindow = 10 * time.Millisecond
	cfg.BackoffMin = 10 * time.Millisecond
	cfg.BackoffMax = 40 * time.Millisecond
	return cfg
}

func eventually(t *testing.T, timeout time.Duration, fn func() bool, msg string) {
	t.Helper()
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if fn() {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatalf("timeout waiting for: %s", msg)
}

func readRender(path string) string {
	b, err := os.ReadFile(path)
	if err != nil {
		return ""
	}
	return string(b)
}

// TestIntegrationFullFlow exercises Enroll → Register → Heartbeat → Subscribe/Ack,
// command application (upsert/revoke), render correctness, heartbeat-triggered
// resync, and reconnect resume via last_command_id.
func TestIntegrationFullFlow(t *testing.T) {
	m := newMockCP(t)
	m.seedCred(&agentv1.Credential{DpUUID: "boot", Protocol: agentv1.ProtocolBoth})
	dial := startMock(t, m)

	cfg := fastConfig(t)
	a := New(cfg,
		WithDialer(Dialer(dial)),
		WithEnrollDialer(EnrollDialer(dial)),
		WithRestarter(&fakeRestarter{}),
	)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	runErr := make(chan error, 1)
	go func() { runErr <- a.Run(ctx) }()

	// 1) Enrollment persisted the cert/key/ca and the node UUID is set.
	eventually(t, 2*time.Second, func() bool {
		_, err1 := os.Stat(cfg.CertPath())
		_, err2 := os.Stat(cfg.KeyPath())
		_, err3 := os.Stat(cfg.CAPath())
		return err1 == nil && err2 == nil && err3 == nil && a.NodeUUID() == testNodeUUID
	}, "enrollment to persist cert/key/ca and set node UUID")

	// 2) Initial Register snapshot rendered with the seeded credential.
	eventually(t, 2*time.Second, func() bool {
		return strings.Contains(readRender(cfg.SingboxConfigPath), `"uuid": "boot"`)
	}, "initial snapshot credential 'boot' to be rendered")

	// 3) Heartbeats flowing.
	eventually(t, 2*time.Second, func() bool { return m.heartbeatCount() > 0 }, "heartbeats")

	// 4) Stream a command: upsert 'extra'. Applied + Ack'd + rendered.
	m.push(&agentv1.Command{
		CommandID: 1, Type: agentv1.CommandTypeUpsert,
		Upsert: &agentv1.UpsertPayload{Credential: &agentv1.Credential{DpUUID: "extra", Protocol: agentv1.ProtocolBoth}},
	})
	eventually(t, 2*time.Second, func() bool {
		return strings.Contains(readRender(cfg.SingboxConfigPath), `"uuid": "extra"`) && m.isAcked(1)
	}, "upsert 'extra' applied and acked")

	// 5) Stream a revoke for 'boot'. Removed from render + Ack'd.
	m.push(&agentv1.Command{
		CommandID: 2, Type: agentv1.CommandTypeRevoke,
		Revoke: &agentv1.RevokePayload{DpUUID: "boot"},
	})
	eventually(t, 2*time.Second, func() bool {
		r := readRender(cfg.SingboxConfigPath)
		return !strings.Contains(r, `"uuid": "boot"`) && m.isAcked(2)
	}, "revoke 'boot' applied and acked")

	// 6) Reconnect resume: enqueue an OFFLINE command (only reachable via backlog),
	//    then drop the stream. After reconnect the agent must resume from
	//    last_command_id and receive id=3 without redelivering 1/2.
	m.enqueueOffline(&agentv1.Command{
		CommandID: 3, Type: agentv1.CommandTypeUpsert,
		Upsert: &agentv1.UpsertPayload{Credential: &agentv1.Credential{DpUUID: "resumed", Protocol: agentv1.ProtocolBoth}},
	})
	m.drop()
	eventually(t, 3*time.Second, func() bool {
		return strings.Contains(readRender(cfg.SingboxConfigPath), `"uuid": "resumed"`) && m.isAcked(3)
	}, "offline command 3 delivered after reconnect resume")

	// The second Subscribe must have resumed from a non-zero high-water mark.
	reqs := m.subReqList()
	if len(reqs) < 2 {
		t.Fatalf("expected at least 2 Subscribe calls (reconnect), got %v", reqs)
	}
	if reqs[len(reqs)-1] < 2 {
		t.Errorf("reconnect Subscribe last_command_id = %d, want >= 2 (resume, no redelivery)", reqs[len(reqs)-1])
	}

	cancel()
	select {
	case <-runErr:
	case <-time.After(2 * time.Second):
		t.Fatal("agent did not shut down after context cancel")
	}
}

// TestIntegrationHeartbeatResync verifies a need_full_resync heartbeat forces a
// reconnect + re-Register so the node converges on the authoritative snapshot.
func TestIntegrationHeartbeatResync(t *testing.T) {
	m := newMockCP(t)
	m.seedCred(&agentv1.Credential{DpUUID: "boot", Protocol: agentv1.ProtocolBoth})
	m.resyncOnce = true
	dial := startMock(t, m)

	a := New(fastConfig(t),
		WithDialer(Dialer(dial)),
		WithEnrollDialer(EnrollDialer(dial)),
		WithRestarter(&fakeRestarter{}),
	)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go func() { _ = a.Run(ctx) }()

	eventually(t, 3*time.Second, func() bool { return m.registerCount() >= 2 }, "re-Register after need_full_resync")
}