DDI/agent/core/server_source_test.go

// SPDX-License-Identifier: Apache-2.0

package core

import (
	"context"
	"errors"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"git.beugsiud.de/beug.k/DDI/internal/agentpki"
	agentv1 "git.beugsiud.de/beug.k/DDI/proto/gen/ddi/agent/v1"
)

// enrollWith enrolls a new store against the fake server.
func enrollWith(t *testing.T, srv *fakeServer, addr string) (*Store, *Identity) {
	t.Helper()
	store := openTestStore(t)
	id, err := Enroll(context.Background(), store, EnrollOptions{
		Kind: srv.role, Server: addr, Token: srv.token, CAHash: srv.fingerprint(),
		ServerName: "localhost", Hostname: "test-host",
	})
	if err != nil {
		t.Fatalf("Enroll: %v", err)
	}
	return store, id
}

func TestEnroll(t *testing.T) {
	srv := newFakeServer(t, KindDNS)
	addr := srv.start(t)
	store, id := enrollWith(t, srv, addr)
	if id.AgentID != srv.agentID || id.Server != addr || id.ServerName != "localhost" || !id.CA.Equal(srv.pki.ca) {
		t.Fatalf("identity = %+v", id)
	}
	loaded, err := store.LoadIdentity()
	if err != nil {
		t.Fatal(err)
	}
	if loaded.AgentID != srv.agentID || !loaded.BundleKey.Equal(srv.signKey.Public()) || loaded.Role != KindDNS {
		t.Errorf("stored identity = %+v", loaded)
	}
	if got, err := agentpki.AgentID(loaded.Certificate.Leaf); err != nil || got != srv.agentID {
		t.Errorf("certificate agent ID = %q, %v", got, err)
	}

	// Enrolling again needs --force.
	opts := EnrollOptions{Kind: KindDNS, Server: addr, Token: srv.token, CAHash: srv.fingerprint(), ServerName: "localhost"}
	if _, err := Enroll(context.Background(), store, opts); !errors.Is(err, errAlreadyEnrolled) {
		t.Fatalf("second Enroll = %v", err)
	}
	srv.with(func(f *fakeServer) { f.agentID = "22222222-2222-4333-8444-555555555555" })
	st := &State{Applied: 3, PendingReports: []DeploymentReport{{Revision: 3}}}
	if err := store.SaveState(st); err != nil {
		t.Fatal(err)
	}
	opts.Force = true
	id2, err := Enroll(context.Background(), store, opts)
	if err != nil || id2.AgentID != "22222222-2222-4333-8444-555555555555" {
		t.Fatalf("forced Enroll = %+v, %v", id2, err)
	}
	if st, _ := store.LoadState(); st.Applied != 3 || len(st.PendingReports) != 0 {
		t.Errorf("state after forced enroll = %+v", st)
	}
}

func TestEnrollDefaultsServerNameToHost(t *testing.T) {
	srv := newFakeServer(t, KindDHCP)
	addr := srv.start(t) // 127.0.0.1:port; the server certificate has the IP SAN
	store := openTestStore(t)
	id, err := Enroll(context.Background(), store, EnrollOptions{
		Kind: KindDHCP, Server: "https://" + addr + "/", Token: srv.token, CAHash: strings.ToUpper(srv.fingerprint()[7:]),
	})
	if err == nil {
		t.Fatalf("fingerprint without prefix accepted: %+v", id)
	}
	id, err = Enroll(context.Background(), store, EnrollOptions{
		Kind: KindDHCP, Server: "https://" + addr + "/", Token: srv.token, CAHash: "sha256:" + strings.ToUpper(srv.fingerprint()[7:]),
	})
	if err != nil {
		t.Fatal(err)
	}
	if id.ServerName != "127.0.0.1" || id.Server != addr {
		t.Errorf("identity = %+v", id)
	}
}

func TestEnrollFailures(t *testing.T) {
	srv := newFakeServer(t, KindDNS)
	addr := srv.start(t)
	other := newTestPKI(t)
	cases := map[string]EnrollOptions{
		"wrong fingerprint": {Server: addr, Token: srv.token, CAHash: agentpki.Fingerprint(other.ca), ServerName: "localhost"},
		"wrong server name": {Server: addr, Token: srv.token, CAHash: srv.fingerprint(), ServerName: "evil.example.com"},
		"wrong token":       {Server: addr, Token: "0123456789abcdef.wrong", CAHash: srv.fingerprint(), ServerName: "localhost"},
		"bad token":         {Server: addr, Token: "nodot", CAHash: srv.fingerprint()},
		"bad address":       {Server: "localhost", Token: srv.token, CAHash: srv.fingerprint()},
		"bad fingerprint":   {Server: addr, Token: srv.token, CAHash: "sha256:abc"},
	}
	for name, opts := range cases {
		t.Run(name, func(t *testing.T) {
			opts.Kind = KindDNS
			store := openTestStore(t)
			if _, err := Enroll(context.Background(), store, opts); err == nil {
				t.Fatal("accepted")
			}
			if store.Enrolled() {
				t.Error("identity stored after a failed enrollment")
			}
		})
	}
	// Role mismatch: the server refuses.
	store := openTestStore(t)
	if _, err := Enroll(context.Background(), store, EnrollOptions{
		Kind: KindDHCP, Server: addr, Token: srv.token, CAHash: srv.fingerprint(), ServerName: "localhost",
	}); err == nil {
		t.Error("wrong role accepted")
	}
	if _, err := Enroll(context.Background(), store, EnrollOptions{
		Kind: Kind("ntp"), Server: addr, Token: srv.token, CAHash: srv.fingerprint(),
	}); err == nil {
		t.Error("unknown role accepted")
	}
	srv.with(func(f *fakeServer) {
		if f.enrollments == 0 {
			t.Error("server never reached")
		}
	})
}

// runAgentAgainst runs an agent with a ServerSource until the test ends.
func runAgentAgainst(t *testing.T, store *Store, id *Identity, kind Kind) (*fakeRole, *ServerSource) {
	t.Helper()
	role := newFakeRole(kind)
	src := NewServerSource(store, id, "test-host", testLogger())
	src.backoffMin, src.backoffMax = 10*time.Millisecond, 100*time.Millisecond
	a, err := NewAgent(kind, testConfig(store.Dir()), role, store, id, src, testLogger())
	if err != nil {
		t.Fatal(err)
	}
	a.healthRetry = 20 * time.Millisecond
	a.reportRetry = 50 * time.Millisecond
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- a.Run(ctx) }()
	t.Cleanup(func() {
		cancel()
		select {
		case err := <-done:
			if err != nil {
				t.Errorf("Run = %v", err)
			}
		case <-time.After(10 * time.Second):
			t.Error("Run did not return after cancel")
		}
	})
	return role, src
}

func serverReports(srv *fakeServer) []*agentv1.ReportDeploymentRequest {
	var out []*agentv1.ReportDeploymentRequest
	srv.with(func(f *fakeServer) { out = append(out, f.reports...) })
	return out
}

func TestServerSourceStreamAndPush(t *testing.T) {
	srv := newFakeServer(t, KindDNS)
	addr := srv.start(t)
	store, id := enrollWith(t, srv, addr)
	srv.with(func(f *fakeServer) { f.desired = 1 })
	role, _ := runAgentAgainst(t, store, id, KindDNS)

	waitFor(t, "report of revision 1", func() bool { return len(serverReports(srv)) == 1 })
	r := serverReports(srv)[0]
	if r.GetRevision() != 1 || r.GetOutcome() != agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_SUCCEEDED ||
		r.GetTrigger() != agentv1.DeploymentTrigger_DEPLOYMENT_TRIGGER_PULL || r.GetActiveRevision() != 1 ||
		r.GetStartedAt() == nil || r.GetFinishedAt() == nil {
		t.Fatalf("report = %v", r)
	}
	srv.with(func(f *fakeServer) {
		if len(f.hellos) == 0 || f.hellos[0].GetHostname() != "test-host" || f.hellos[0].GetAgentVersion() == "" {
			t.Errorf("hellos = %v", f.hellos)
		}
	})
	waitFor(t, "heartbeats", func() bool {
		n := 0
		srv.with(func(f *fakeServer) { n = f.heartbeats })
		return n >= 2
	})

	// Push over the stream.
	srv.commands <- &agentv1.Command{CommandId: "c1", Action: &agentv1.Command_ApplyRevision{
		ApplyRevision: &agentv1.ApplyRevisionCommand{Revision: 2, DeploymentId: "dep-2"},
	}}
	waitFor(t, "report of revision 2", func() bool { return len(serverReports(srv)) == 2 })
	r = serverReports(srv)[1]
	if r.GetRevision() != 2 || r.GetDeploymentId() != "dep-2" || r.GetTrigger() != agentv1.DeploymentTrigger_DEPLOYMENT_TRIGGER_PUSH ||
		r.GetOutcome() != agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_SUCCEEDED {
		t.Fatalf("report = %v", r)
	}
	srv.with(func(f *fakeServer) {
		if len(f.acks) != 1 || f.acks[0].GetCommandId() != "c1" || !f.acks[0].GetAccepted() {
			t.Errorf("acks = %v", f.acks)
		}
	})
	if n := role.activeRevision(); n != 2 {
		t.Errorf("active = %d", n)
	}

	// Invalid and unknown commands are refused.
	srv.commands <- &agentv1.Command{CommandId: "c2", Action: &agentv1.Command_ApplyRevision{ApplyRevision: &agentv1.ApplyRevisionCommand{}}}
	srv.commands <- &agentv1.Command{CommandId: "c3"}
	waitFor(t, "refusals", func() bool {
		n := 0
		srv.with(func(f *fakeServer) { n = len(f.acks) })
		return n == 3
	})
	srv.with(func(f *fakeServer) {
		if f.acks[1].GetAccepted() || f.acks[2].GetAccepted() {
			t.Errorf("acks = %v", f.acks)
		}
	})
}

func TestServerSourceManualMode(t *testing.T) {
	srv := newFakeServer(t, KindDNS)
	addr := srv.start(t)
	store, id := enrollWith(t, srv, addr)
	srv.with(func(f *fakeServer) {
		f.desired = 1
		f.mode = agentv1.OperatingMode_OPERATING_MODE_MANUAL
	})
	runAgentAgainst(t, store, id, KindDNS)
	waitFor(t, "hello", func() bool {
		n := 0
		srv.with(func(f *fakeServer) { n = f.heartbeats })
		return n >= 3
	})
	if n := len(serverReports(srv)); n != 0 {
		t.Errorf("manual mode applied the desired revision (%d reports)", n)
	}
}

func TestServerSourcePollingFallback(t *testing.T) {
	srv := newFakeServer(t, KindDNS)
	addr := srv.start(t)
	store, id := enrollWith(t, srv, addr)
	srv.with(func(f *fakeServer) {
		f.desired = 1
		f.streamDisabled = true
	})
	runAgentAgainst(t, store, id, KindDNS)
	waitFor(t, "report via polling", func() bool { return len(serverReports(srv)) == 1 })
	srv.with(func(f *fakeServer) {
		if f.polls == 0 {
			t.Error("GetDesiredState never called")
		}
	})
	// A newer revision is picked up by a later poll.
	srv.with(func(f *fakeServer) { f.desired = 2 })
	waitFor(t, "report of revision 2", func() bool { return len(serverReports(srv)) == 2 })
	if st, _ := store.LoadState(); st.Applied != 2 {
		t.Errorf("state = %+v", st)
	}
}

func TestServerSourceRenewsCertificate(t *testing.T) {
	srv := newFakeServer(t, KindDNS)
	// More than half of the certificate's lifetime has passed.
	srv.enrollValidity = [2]time.Duration{-3 * time.Hour, time.Hour}
	addr := srv.start(t)
	store, id := enrollWith(t, srv, addr)
	oldCert, _ := os.ReadFile(filepath.Join(store.Dir(), "identity", "cert.pem"))
	runAgentAgainst(t, store, id, KindDNS)

	waitFor(t, "renewal and hello", func() bool {
		ok := false
		srv.with(func(f *fakeServer) { ok = f.renewals == 1 && len(f.hellos) >= 1 })
		return ok
	})
	newCert, _ := os.ReadFile(filepath.Join(store.Dir(), "identity", "cert.pem"))
	if string(newCert) == string(oldCert) {
		t.Error("cert.pem unchanged after renewal")
	}
	loaded, err := store.LoadIdentity()
	if err != nil {
		t.Fatal(err)
	}
	if time.Until(loaded.Certificate.Leaf.NotAfter) < 24*time.Hour {
		t.Errorf("renewed certificate expires %v", loaded.Certificate.Leaf.NotAfter)
	}
}

func TestServerSourceRenewCommand(t *testing.T) {
	srv := newFakeServer(t, KindDNS)
	addr := srv.start(t)
	store, id := enrollWith(t, srv, addr)
	runAgentAgainst(t, store, id, KindDNS)
	waitFor(t, "first hello", func() bool {
		n := 0
		srv.with(func(f *fakeServer) { n = len(f.hellos) })
		return n == 1
	})
	srv.commands <- &agentv1.Command{CommandId: "r1", Action: &agentv1.Command_RenewCertificate{RenewCertificate: &agentv1.RenewCertificateCommand{}}}
	// The old certificate is no longer accepted, so the second Hello proves
	// that the agent reconnected with the new one.
	waitFor(t, "reconnect with the renewed certificate", func() bool {
		ok := false
		srv.with(func(f *fakeServer) { ok = f.renewals == 1 && len(f.hellos) >= 2 })
		return ok
	})
}

func TestServerSourceReportsLeases(t *testing.T) {
	srv := newFakeServer(t, KindDHCP)
	addr := srv.start(t)
	store, id := enrollWith(t, srv, addr)
	role, _ := runAgentAgainst(t, store, id, KindDHCP)
	_ = role
	waitFor(t, "full lease sync", func() bool {
		ok := false
		srv.with(func(f *fakeServer) {
			for _, r := range f.leaseReports {
				if r.GetFullSync() {
					ok = true
				}
			}
		})
		return ok
	})
}

func TestLeaseToProto(t *testing.T) {
	l := lease("10.0.0.9", "aa:bb:cc:dd:ee:ff")
	l.ClientID = "01:aa"
	p := leaseToProto(l)
	if p.GetIp() != "10.0.0.9" || p.GetMac() != l.MAC || p.GetState() != agentv1.LeaseState_LEASE_STATE_ACTIVE ||
		!p.GetExpiresAt().AsTime().Equal(l.Expires) || p.GetSubnetId() != 1 || p.GetClientId() != "01:aa" {
		t.Errorf("lease = %v", p)
	}
	if p := leaseToProto(Lease{IP: mustAddr("10.0.0.1")}); p.ExpiresAt != nil || p.UpdatedAt != nil {
		t.Errorf("zero times encoded: %v", p)
	}
	for s, want := range map[LeaseState]agentv1.LeaseState{
		LeaseUnknown: agentv1.LeaseState_LEASE_STATE_UNSPECIFIED, LeaseExpired: agentv1.LeaseState_LEASE_STATE_EXPIRED,
		LeaseReleased: agentv1.LeaseState_LEASE_STATE_RELEASED, LeaseDeclined: agentv1.LeaseState_LEASE_STATE_DECLINED,
	} {
		if got := leaseStateToProto(s); got != want {
			t.Errorf("%v -> %v, want %v", s, got, want)
		}
	}
	for o, want := range map[Outcome]agentv1.DeploymentOutcome{
		OutcomeRejected: agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_REJECTED, OutcomeRolledBack: agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_ROLLED_BACK,
		OutcomeFailed: agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_FAILED, Outcome("x"): agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_UNSPECIFIED,
	} {
		if got := outcomeToProto(o); got != want {
			t.Errorf("%v -> %v, want %v", o, got, want)
		}
	}
	if triggerToProto(Trigger("x")) != agentv1.DeploymentTrigger_DEPLOYMENT_TRIGGER_UNSPECIFIED {
		t.Error("unknown trigger mapped")
	}
	h := healthToProto(Health{State: HealthDegraded, Checks: []HealthCheck{{Name: "a", State: HealthFailed}}}, time.Time{})
	if h.GetState() != agentv1.HealthState_HEALTH_STATE_DEGRADED || h.CheckedAt != nil || h.GetChecks()[0].GetState() != agentv1.HealthState_HEALTH_STATE_FAILED {
		t.Errorf("health = %v", h)
	}
}

func TestServerSourceUnreachable(t *testing.T) {
	srv := newFakeServer(t, KindDNS)
	addr := srv.start(t)
	store, id := enrollWith(t, srv, addr)
	id.Server = "127.0.0.1:1" // nothing listens there
	src := NewServerSource(store, id, "h", testLogger())
	defer src.closeConn()
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	if err := src.ReportDeployment(ctx, DeploymentReport{Revision: 1}); err == nil {
		t.Error("report to an unreachable server succeeded")
	}
	if _, err := src.FetchBundle(ctx, 1); err == nil {
		t.Error("fetch from an unreachable server succeeded")
	}
	if code := status.Code(src.ReportLeases(ctx, LeaseBatch{Sequence: 1})); code == codes.OK {
		t.Error("lease report to an unreachable server succeeded")
	}
}

func TestJitterAndRenewalDue(t *testing.T) {
	for i := 0; i < 100; i++ {
		if d := jitter(time.Second); d < 800*time.Millisecond || d > 1200*time.Millisecond {
			t.Fatalf("jitter = %v", d)
		}
	}
	pki := newTestPKI(t)
	id, _ := testIdentity(t, pki)
	leaf := id.Certificate.Leaf
	if due := renewalDue(leaf); !due.Equal(leaf.NotBefore.Add(leaf.NotAfter.Sub(leaf.NotBefore) / 2)) {
		t.Errorf("renewalDue = %v", due)
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if sleepCtx(ctx, time.Hour) {
		t.Error("sleepCtx ignored the context")
	}
}