DDI/agent/core/helpers_test.go

// SPDX-License-Identifier: Apache-2.0

package core

import (
	"context"
	"crypto"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"math/big"
	"net"
	"net/url"
	"os"
	"path/filepath"
	"slices"
	"sync"
	"testing"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/peer"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/durationpb"

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

const testAgentID = "0b6f3c1e-5d2a-4c8e-9f10-123456789abc"

func testLogger() *slog.Logger {
	if testing.Verbose() {
		return slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelDebug}))
	}
	return slog.New(slog.NewTextHandler(io.Discard, nil))
}

// waitFor polls cond until it holds or the timeout expires.
func waitFor(t *testing.T, what string, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(15 * time.Second)
	for !cond() {
		if time.Now().After(deadline) {
			t.Fatalf("timeout waiting for %s", what)
		}
		time.Sleep(10 * time.Millisecond)
	}
}

// ---- PKI ----

type testPKI struct {
	ca         *x509.Certificate
	caKey      *ecdsa.PrivateKey
	serverCert tls.Certificate
}

func newTestPKI(t *testing.T) *testPKI {
	t.Helper()
	caKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	tmpl := &x509.Certificate{
		SerialNumber: big.NewInt(1), Subject: pkix.Name{CommonName: "test agent CA"},
		NotBefore: time.Now().Add(-24 * time.Hour), NotAfter: time.Now().Add(24 * time.Hour),
		IsCA: true, BasicConstraintsValid: true, KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature,
	}
	der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &caKey.PublicKey, caKey)
	if err != nil {
		t.Fatal(err)
	}
	ca, err := x509.ParseCertificate(der)
	if err != nil {
		t.Fatal(err)
	}
	p := &testPKI{ca: ca, caKey: caKey}
	srvKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	srvDER, err := x509.CreateCertificate(rand.Reader, &x509.Certificate{
		SerialNumber: big.NewInt(2), Subject: pkix.Name{CommonName: "localhost"},
		DNSNames: []string{"localhost"}, IPAddresses: []net.IP{net.IPv4(127, 0, 0, 1)},
		NotBefore: time.Now().Add(-time.Hour), NotAfter: time.Now().Add(24 * time.Hour),
		KeyUsage: x509.KeyUsageDigitalSignature, ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}, ca, &srvKey.PublicKey, caKey)
	if err != nil {
		t.Fatal(err)
	}
	p.serverCert = tls.Certificate{Certificate: [][]byte{srvDER, ca.Raw}, PrivateKey: srvKey}
	return p
}

func (p *testPKI) issueClient(pub crypto.PublicKey, agentID string, notBefore, notAfter time.Time) (*x509.Certificate, error) {
	serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 127))
	if err != nil {
		return nil, err
	}
	der, err := x509.CreateCertificate(rand.Reader, &x509.Certificate{
		SerialNumber: serial, Subject: pkix.Name{CommonName: agentID},
		URIs:      []*url.URL{agentpki.AgentURI(agentID)},
		NotBefore: notBefore, NotAfter: notAfter,
		KeyUsage: x509.KeyUsageDigitalSignature, ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, p.ca, pub, p.caKey)
	if err != nil {
		return nil, err
	}
	return x509.ParseCertificate(der)
}

// ---- bundles ----

func newSigningKey(t *testing.T) (ed25519.PublicKey, ed25519.PrivateKey) {
	t.Helper()
	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	return pub, priv
}

// testArtifacts returns the artifacts of a test revision.
func testArtifacts(rev uint64) []bundle.Artifact {
	return []bundle.Artifact{
		{Path: "ddi-suite.conf", Mode: 0o640, Content: fmt.Appendf(nil, "// revision %d\n", rev)},
		{Path: "zones/db.example.com", Mode: 0o644, Content: fmt.Appendf(nil, "example.com. SOA ns 1 %d\n", rev)},
	}
}

func signBundle(t testing.TB, key ed25519.PrivateKey, agentID string, role Kind, rev uint64) *Bundle {
	t.Helper()
	arts := testArtifacts(rev)
	manifest, sig, err := bundle.Sign(key, agentID, string(role), rev, time.Unix(1_700_000_000, 0), arts)
	if err != nil {
		t.Fatal(err)
	}
	b := &Bundle{Manifest: manifest, Signature: sig}
	for _, a := range arts {
		b.Artifacts = append(b.Artifacts, BundleArtifact{Path: a.Path, Content: a.Content})
	}
	return b
}

// ---- fake role ----

type fakeRole struct {
	kind Kind

	mu          sync.Mutex
	validateErr map[uint64]error
	activateErr map[uint64]error
	unhealthy   map[uint64]bool
	active      uint64
	calls       []string
	healthCalls map[uint64]int
	leases      []Lease
	leasesErr   error
	versionErr  error
}

func newFakeRole(kind Kind) *fakeRole {
	return &fakeRole{
		kind: kind, validateErr: map[uint64]error{}, activateErr: map[uint64]error{},
		unhealthy: map[uint64]bool{}, healthCalls: map[uint64]int{},
	}
}

func (r *fakeRole) Kind() Kind { return r.kind }

func (r *fakeRole) ServiceVersion(context.Context) (string, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.versionErr != nil {
		return "", r.versionErr
	}
	return "Fake 1.0", nil
}

// checkFiles verifies that the stored revision holds its artifacts.
func checkFiles(rev Revision) error {
	for _, f := range rev.Files {
		data, err := os.ReadFile(filepath.Join(rev.Dir, filepath.FromSlash(f.Path)))
		if err != nil {
			return err
		}
		if int64(len(data)) != f.Size {
			return fmt.Errorf("%s: size %d, want %d", f.Path, len(data), f.Size)
		}
	}
	return nil
}

func (r *fakeRole) Validate(_ context.Context, rev Revision) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.calls = append(r.calls, fmt.Sprintf("validate %d", rev.Number))
	if err := checkFiles(rev); err != nil {
		return err
	}
	return r.validateErr[rev.Number]
}

func (r *fakeRole) Activate(_ context.Context, rev Revision, prev *Revision) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	p := "none"
	if prev != nil {
		p = fmt.Sprint(prev.Number)
	}
	r.calls = append(r.calls, fmt.Sprintf("activate %d prev %s", rev.Number, p))
	if err := r.activateErr[rev.Number]; err != nil {
		return err
	}
	r.active = rev.Number
	return nil
}

func (r *fakeRole) Health(_ context.Context, expected *Revision) Health {
	r.mu.Lock()
	defer r.mu.Unlock()
	var n uint64
	if expected != nil {
		n = expected.Number
	}
	r.healthCalls[n]++
	if r.unhealthy[r.active] {
		return Health{State: HealthFailed, Message: "service down",
			Checks: []HealthCheck{{Name: "soa", State: HealthFailed, Message: "serial mismatch"}, {Name: "status", State: HealthOK}}}
	}
	return Health{State: HealthOK}
}

func (r *fakeRole) Leases(context.Context) ([]Lease, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	return slices.Clone(r.leases), r.leasesErr
}

func (r *fakeRole) setLeases(l []Lease) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.leases = l
}

func (r *fakeRole) callLog() []string {
	r.mu.Lock()
	defer r.mu.Unlock()
	return slices.Clone(r.calls)
}

func (r *fakeRole) activeRevision() uint64 {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.active
}

func (r *fakeRole) resetCalls() {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.calls = nil
}

// fakeLeaseRole is a role with a LeaseSource; fakeRole itself has the
// method too, so wrap it to hide Leases for the DNS case.
type dnsOnlyRole struct{ *fakeRole }

func (d dnsOnlyRole) Leases() {} // shadows LeaseSource.Leases with another signature

// ---- fake source ----

type fakeSource struct {
	mu        sync.Mutex
	bundles   map[uint64]*Bundle
	fetchErr  error
	reportErr error
	leaseErr  error
	reports   []DeploymentReport
	batches   []LeaseBatch
	fetches   int
	onFetch   func()
	run       func(ctx context.Context, h SourceHandler) error
}

func newFakeSource() *fakeSource { return &fakeSource{bundles: map[uint64]*Bundle{}} }

func (s *fakeSource) Run(ctx context.Context, h SourceHandler) error {
	if s.run != nil {
		return s.run(ctx, h)
	}
	<-ctx.Done()
	return ctx.Err()
}

func (s *fakeSource) FetchBundle(_ context.Context, rev uint64) (*Bundle, error) {
	s.mu.Lock()
	s.fetches++
	b, err, hook := s.bundles[rev], s.fetchErr, s.onFetch
	s.mu.Unlock()
	if hook != nil {
		hook()
	}
	if err != nil {
		return nil, err
	}
	if b == nil {
		return nil, fmt.Errorf("revision %d not found", rev)
	}
	return b, nil
}

func (s *fakeSource) ReportDeployment(_ context.Context, r DeploymentReport) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.reportErr != nil {
		return s.reportErr
	}
	s.reports = append(s.reports, r)
	return nil
}

func (s *fakeSource) ReportLeases(_ context.Context, b LeaseBatch) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.leaseErr != nil {
		return s.leaseErr
	}
	b.Leases = slices.Clone(b.Leases)
	s.batches = append(s.batches, b)
	return nil
}

func (s *fakeSource) setReportErr(err error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.reportErr = err
}

func (s *fakeSource) reported() []DeploymentReport {
	s.mu.Lock()
	defer s.mu.Unlock()
	return slices.Clone(s.reports)
}

func (s *fakeSource) leaseBatches() []LeaseBatch {
	s.mu.Lock()
	defer s.mu.Unlock()
	return slices.Clone(s.batches)
}

// ---- agent fixture ----

type fixture struct {
	agent  *Agent
	role   *fakeRole
	src    *fakeSource
	store  *Store
	signer ed25519.PrivateKey
}

func testConfig(dir string) *Config {
	return &Config{
		StateDir: dir, ApplyTimeout: 5 * time.Second, HealthTimeout: 300 * time.Millisecond,
		LeaseInterval: time.Second, LeaseFullSync: time.Hour,
	}
}

func newFixture(t *testing.T, kind Kind) *fixture {
	t.Helper()
	store, err := OpenStore(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	return newFixtureWithStore(t, kind, store, nil)
}

func newFixtureWithStore(t *testing.T, kind Kind, store *Store, signer ed25519.PrivateKey) *fixture {
	t.Helper()
	if signer == nil {
		_, signer = newSigningKey(t)
	}
	role := newFakeRole(kind)
	src := newFakeSource()
	id := &Identity{AgentID: testAgentID, Role: kind, BundleKey: signer.Public().(ed25519.PublicKey)}
	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
	return &fixture{agent: a, role: role, src: src, store: store, signer: signer}
}

// addBundle makes revision rev available from the fake source.
func (f *fixture) addBundle(t *testing.T, rev uint64) {
	t.Helper()
	f.src.mu.Lock()
	defer f.src.mu.Unlock()
	f.src.bundles[rev] = signBundle(t, f.signer, testAgentID, f.agent.kind, rev)
}

func (f *fixture) pull(rev uint64) {
	f.agent.handle(context.Background(), applyRequest{revision: rev, trigger: TriggerPull})
}

func (f *fixture) push(rev uint64, dep string) {
	f.agent.handle(context.Background(), applyRequest{revision: rev, trigger: TriggerPush, deploymentID: dep})
}

func (f *fixture) lastResult(t *testing.T) DeploymentReport {
	t.Helper()
	st := f.agent.state()
	if st.LastResult == nil {
		t.Fatal("no result recorded")
	}
	return *st.LastResult
}

// ---- fake gRPC server ----

type fakeServer struct {
	agentv1.UnimplementedAgentServiceServer
	pki     *testPKI
	signKey ed25519.PrivateKey
	role    Kind

	mu             sync.Mutex
	token          string
	agentID        string
	serial         *big.Int
	enrollValidity [2]time.Duration // NotBefore offset, NotAfter offset for Enroll
	desired        uint64
	mode           agentv1.OperatingMode
	pullInterval   time.Duration
	heartbeat      time.Duration
	streamDisabled bool
	hellos         []*agentv1.Hello
	heartbeats     int
	acks           []*agentv1.CommandAck
	reports        []*agentv1.ReportDeploymentRequest
	leaseReports   []*agentv1.ReportLeasesRequest
	renewals       int
	polls          int
	commands       chan *agentv1.Command
	enrollments    int
}

func newFakeServer(t *testing.T, role Kind) *fakeServer {
	t.Helper()
	_, key := newSigningKey(t)
	return &fakeServer{
		pki: newTestPKI(t), signKey: key, role: role, token: "0123456789abcdef.secret",
		agentID:        "11111111-2222-4333-8444-555555555555",
		enrollValidity: [2]time.Duration{-time.Minute, 14 * 24 * time.Hour},
		mode:           agentv1.OperatingMode_OPERATING_MODE_AUTO_PULL,
		pullInterval:   50 * time.Millisecond, heartbeat: 50 * time.Millisecond,
		commands: make(chan *agentv1.Command, 10),
	}
}

// start serves the fake on 127.0.0.1 with mTLS and returns its address.
func (f *fakeServer) start(t *testing.T) string {
	t.Helper()
	lis, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	pool := x509.NewCertPool()
	pool.AddCert(f.pki.ca)
	srv := grpc.NewServer(
		grpc.Creds(credentials.NewTLS(&tls.Config{
			MinVersion: tls.VersionTLS13, Certificates: []tls.Certificate{f.pki.serverCert},
			ClientAuth: tls.VerifyClientCertIfGiven, ClientCAs: pool,
		})),
		grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{MinTime: 10 * time.Second, PermitWithoutStream: true}),
	)
	agentv1.RegisterAgentServiceServer(srv, f)
	go func() { _ = srv.Serve(lis) }()
	t.Cleanup(srv.Stop)
	return lis.Addr().String()
}

func (f *fakeServer) fingerprint() string { return agentpki.Fingerprint(f.pki.ca) }

func (f *fakeServer) authenticate(ctx context.Context) error {
	p, ok := peer.FromContext(ctx)
	if !ok {
		return status.Error(codes.Unauthenticated, "no peer")
	}
	ti, ok := p.AuthInfo.(credentials.TLSInfo)
	if !ok || len(ti.State.PeerCertificates) == 0 {
		return status.Error(codes.Unauthenticated, "client certificate required")
	}
	cert := ti.State.PeerCertificates[0]
	id, err := agentpki.AgentID(cert)
	if err != nil {
		return status.Error(codes.Unauthenticated, err.Error())
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	if id != f.agentID || f.serial == nil || cert.SerialNumber.Cmp(f.serial) != 0 {
		return status.Error(codes.Unauthenticated, "unknown certificate")
	}
	return nil
}

func (f *fakeServer) issue(csrDER []byte, notBefore, notAfter time.Time) (*x509.Certificate, error) {
	csr, err := x509.ParseCertificateRequest(csrDER)
	if err != nil {
		return nil, err
	}
	if err := csr.CheckSignature(); err != nil {
		return nil, err
	}
	f.mu.Lock()
	id := f.agentID
	f.mu.Unlock()
	cert, err := f.pki.issueClient(csr.PublicKey, id, notBefore, notAfter)
	if err != nil {
		return nil, err
	}
	f.mu.Lock()
	f.serial = cert.SerialNumber
	f.mu.Unlock()
	return cert, nil
}

func (f *fakeServer) Enroll(_ context.Context, req *agentv1.EnrollRequest) (*agentv1.EnrollResponse, error) {
	f.mu.Lock()
	token, validity := f.token, f.enrollValidity
	f.enrollments++
	f.mu.Unlock()
	if req.GetToken() != token {
		return nil, status.Error(codes.PermissionDenied, "invalid enrollment token")
	}
	want, _ := roleToProto(f.role)
	if req.GetRole() != want {
		return nil, status.Error(codes.InvalidArgument, "wrong role")
	}
	now := time.Now()
	cert, err := f.issue(req.GetCsrDer(), now.Add(validity[0]), now.Add(validity[1]))
	if err != nil {
		return nil, status.Error(codes.InvalidArgument, err.Error())
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	return &agentv1.EnrollResponse{
		AgentId: f.agentID, CertificateDer: cert.Raw, CaCertificatesDer: [][]byte{f.pki.ca.Raw},
		BundleSigningPublicKey: f.signKey.Public().(ed25519.PublicKey),
	}, nil
}

func (f *fakeServer) RenewCertificate(ctx context.Context, req *agentv1.RenewCertificateRequest) (*agentv1.RenewCertificateResponse, error) {
	if err := f.authenticate(ctx); err != nil {
		return nil, err
	}
	now := time.Now()
	cert, err := f.issue(req.GetCsrDer(), now.Add(-time.Minute), now.Add(14*24*time.Hour))
	if err != nil {
		return nil, status.Error(codes.InvalidArgument, err.Error())
	}
	f.mu.Lock()
	f.renewals++
	f.mu.Unlock()
	return &agentv1.RenewCertificateResponse{CertificateDer: cert.Raw, CaCertificatesDer: [][]byte{f.pki.ca.Raw}}, nil
}

func (f *fakeServer) Connect(stream agentv1.AgentService_ConnectServer) error {
	f.mu.Lock()
	disabled := f.streamDisabled
	f.mu.Unlock()
	if disabled {
		return status.Error(codes.Unimplemented, "streams disabled")
	}
	if err := f.authenticate(stream.Context()); err != nil {
		return err
	}
	first, err := stream.Recv()
	if err != nil {
		return err
	}
	if first.GetHello() == nil {
		return status.Error(codes.InvalidArgument, "Hello expected")
	}
	f.mu.Lock()
	f.hellos = append(f.hellos, first.GetHello())
	w := &agentv1.Welcome{
		Mode: f.mode, PullInterval: durationpb.New(f.pullInterval),
		HeartbeatInterval: durationpb.New(f.heartbeat), DesiredRevision: f.desired,
	}
	f.mu.Unlock()
	if err := stream.Send(&agentv1.ConnectResponse{Message: &agentv1.ConnectResponse_Welcome{Welcome: w}}); err != nil {
		return err
	}
	recvErr := make(chan error, 1)
	go func() {
		for {
			msg, err := stream.Recv()
			if err != nil {
				recvErr <- err
				return
			}
			f.mu.Lock()
			switch {
			case msg.GetHeartbeat() != nil:
				f.heartbeats++
			case msg.GetCommandAck() != nil:
				f.acks = append(f.acks, msg.GetCommandAck())
			}
			f.mu.Unlock()
		}
	}()
	for {
		select {
		case err := <-recvErr:
			if errors.Is(err, io.EOF) {
				return nil
			}
			return err
		case <-stream.Context().Done():
			return nil
		case c := <-f.commands:
			if err := stream.Send(&agentv1.ConnectResponse{Message: &agentv1.ConnectResponse_Command{Command: c}}); err != nil {
				return err
			}
		}
	}
}

func (f *fakeServer) GetDesiredState(ctx context.Context, _ *agentv1.GetDesiredStateRequest) (*agentv1.GetDesiredStateResponse, error) {
	if err := f.authenticate(ctx); err != nil {
		return nil, err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	f.polls++
	return &agentv1.GetDesiredStateResponse{DesiredRevision: f.desired, Mode: f.mode, PullInterval: durationpb.New(f.pullInterval)}, nil
}

func (f *fakeServer) FetchBundle(ctx context.Context, req *agentv1.FetchBundleRequest) (*agentv1.FetchBundleResponse, error) {
	if err := f.authenticate(ctx); err != nil {
		return nil, err
	}
	f.mu.Lock()
	id := f.agentID
	f.mu.Unlock()
	arts := testArtifacts(req.GetRevision())
	manifest, sig, err := bundle.Sign(f.signKey, id, string(f.role), req.GetRevision(), time.Unix(1_700_000_000, 0), arts)
	if err != nil {
		return nil, status.Error(codes.Internal, err.Error())
	}
	resp := &agentv1.FetchBundleResponse{Manifest: manifest, Signature: sig}
	for _, a := range arts {
		resp.Artifacts = append(resp.Artifacts, &agentv1.Artifact{Path: a.Path, Content: a.Content})
	}
	return resp, nil
}

func (f *fakeServer) ReportDeployment(ctx context.Context, req *agentv1.ReportDeploymentRequest) (*agentv1.ReportDeploymentResponse, error) {
	if err := f.authenticate(ctx); err != nil {
		return nil, err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	f.reports = append(f.reports, req)
	return &agentv1.ReportDeploymentResponse{}, nil
}

func (f *fakeServer) ReportLeases(ctx context.Context, req *agentv1.ReportLeasesRequest) (*agentv1.ReportLeasesResponse, error) {
	if err := f.authenticate(ctx); err != nil {
		return nil, err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	f.leaseReports = append(f.leaseReports, req)
	return &agentv1.ReportLeasesResponse{AcknowledgedSequence: req.GetSequence()}, nil
}

// snapshot returns a copy of the recorded data under the lock.
func (f *fakeServer) with(fn func(f *fakeServer)) {
	f.mu.Lock()
	defer f.mu.Unlock()
	fn(f)
}