DDI/server/internal/agentapi/server_test.go

// SPDX-License-Identifier: LicenseRef-DDI-Proprietary

package agentapi

import (
	"bytes"
	"context"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"io"
	"log/slog"
	"net"
	"sync/atomic"
	"testing"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/timestamppb"

	"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"
	"git.beugsiud.de/beug.k/DDI/server/internal/agenthub"
	"git.beugsiud.de/beug.k/DDI/server/internal/pki"
	"git.beugsiud.de/beug.k/DDI/server/internal/secrets"
	"git.beugsiud.de/beug.k/DDI/server/internal/service"
	"git.beugsiud.de/beug.k/DDI/server/internal/store"
	"git.beugsiud.de/beug.k/DDI/server/internal/store/storetest"
)

const serverName = "agents.test"

type testEnv struct {
	t   *testing.T
	st  *store.Store
	svc *service.Service
	pki *pki.Authority
	hub *agenthub.Hub
	srv *Server
	// offset shifts the PKI clock (certificate issuance and checks).
	offset atomic.Int64
	addr   string
}

func newTestEnv(t *testing.T) *testEnv {
	t.Helper()
	e := &testEnv{t: t, st: storetest.New(t)}
	keyData, _ := secrets.GenerateKeyFile("k1")
	kr, err := secrets.ParseKeyFile(keyData)
	if err != nil {
		t.Fatal(err)
	}
	clock := func() time.Time { return time.Now().Add(time.Duration(e.offset.Load())) }
	// The CA must predate the backdated agent certificates of the renewal
	// tests (chains are checked at the time a certificate was valid).
	e.offset.Store(int64(-90 * 24 * time.Hour))
	e.pki, err = pki.LoadOrCreate(context.Background(), e.st, kr, nil, clock)
	e.offset.Store(0)
	if err != nil {
		t.Fatal(err)
	}
	e.hub = agenthub.New()
	log := slog.New(slog.DiscardHandler)
	e.svc = service.New(service.Options{Store: e.st, Keyring: kr, PKI: e.pki, Hub: e.hub, Logger: log,
		AgentAPIAddress: serverName + ":8444", AgentServerNames: []string{serverName}})
	e.srv, err = New(Options{Service: e.svc, PKI: e.pki, Hub: e.hub, ServerNames: []string{serverName}, Logger: log})
	if err != nil {
		t.Fatal(err)
	}
	lis, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	e.addr = lis.Addr().String()
	done := make(chan error, 1)
	go func() { done <- e.srv.Serve(lis) }()
	t.Cleanup(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		e.srv.Shutdown(ctx)
		if err := <-done; err != nil {
			t.Errorf("serve: %v", err)
		}
	})
	return e
}

func (e *testEnv) dial(cfg *tls.Config) agentv1.AgentServiceClient {
	e.t.Helper()
	cc, err := grpc.NewClient("passthrough:///"+e.addr, grpc.WithTransportCredentials(credentials.NewTLS(cfg)))
	if err != nil {
		e.t.Fatal(err)
	}
	e.t.Cleanup(func() { _ = cc.Close() })
	return agentv1.NewAgentServiceClient(cc)
}

// bootstrapClient trusts the server only through the CA fingerprint, like
// the agent's enroll command.
func (e *testEnv) bootstrapClient(fingerprint string) agentv1.AgentServiceClient {
	return e.dial(&tls.Config{
		MinVersion: tls.VersionTLS13, ServerName: serverName,
		InsecureSkipVerify: true, //nolint:gosec // verified by VerifyPeerCertificate
		VerifyPeerCertificate: func(raw [][]byte, _ [][]*x509.Certificate) error {
			_, err := agentpki.VerifyChainWithFingerprint(raw, fingerprint, serverName)
			return err
		},
	})
}

func (e *testEnv) token(role string) string {
	e.t.Helper()
	tk, err := e.svc.CreateEnrollmentToken(context.Background(), service.SystemActor, service.TokenInput{Role: role})
	if err != nil {
		e.t.Fatal(err)
	}
	return tk.Secret
}

func csr(t *testing.T) (*ecdsa.PrivateKey, []byte) {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	der, err := x509.CreateCertificateRequest(rand.Reader, &x509.CertificateRequest{}, key)
	if err != nil {
		t.Fatal(err)
	}
	return key, der
}

type testAgent struct {
	id        string
	role      string
	ca        *x509.Certificate
	cert      *x509.Certificate
	bundlePub ed25519.PublicKey
	client    agentv1.AgentServiceClient
}

func ctxT(t *testing.T) context.Context {
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	t.Cleanup(cancel)
	return ctx
}

func roleProto(role string) agentv1.Role {
	if role == "dhcp" {
		return agentv1.Role_ROLE_DHCP
	}
	return agentv1.Role_ROLE_DNS
}

func (e *testEnv) enroll(role string) *testAgent {
	e.t.Helper()
	key, req := csr(e.t)
	resp, err := e.bootstrapClient(e.pki.Fingerprint()).Enroll(ctxT(e.t), &agentv1.EnrollRequest{
		Token: e.token(role), CsrDer: req, Role: roleProto(role), Hostname: "host-" + role, AgentVersion: "1.0.0",
	})
	if err != nil {
		e.t.Fatalf("enroll: %v", err)
	}
	return e.agentFrom(resp.GetAgentId(), role, key, resp.GetCertificateDer(), resp.GetCaCertificatesDer(), resp.GetBundleSigningPublicKey())
}

func (e *testEnv) agentFrom(id, role string, key *ecdsa.PrivateKey, certDER []byte, chain [][]byte, pub []byte) *testAgent {
	e.t.Helper()
	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		e.t.Fatal(err)
	}
	if len(chain) != 1 {
		e.t.Fatalf("CA chain length %d", len(chain))
	}
	ca, err := x509.ParseCertificate(chain[len(chain)-1])
	if err != nil {
		e.t.Fatal(err)
	}
	client := e.dial(agentpki.ClientTLSConfig(ca, tls.Certificate{Certificate: [][]byte{certDER}, PrivateKey: key}, serverName))
	return &testAgent{id: id, role: role, ca: ca, cert: cert, bundlePub: pub, client: client}
}

func wantCode(t *testing.T, err error, want codes.Code) {
	t.Helper()
	if status.Code(err) != want {
		t.Errorf("code %v (%v), want %v", status.Code(err), err, want)
	}
}

// eventually polls cond for up to 5 s.
func eventually(t *testing.T, what string, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(5 * time.Second)
	for !cond() {
		if time.Now().After(deadline) {
			t.Fatalf("timed out waiting for %s", what)
		}
		time.Sleep(20 * time.Millisecond)
	}
}

func (e *testEnv) configureGenerator(ttl uint32) {
	e.t.Helper()
	ctx := context.Background()
	if _, err := e.st.Pool().Exec(ctx, `INSERT INTO networks (cidr, name, domain) VALUES ('192.0.2.0/24', 'Office', 'office.test')
		ON CONFLICT DO NOTHING`); err != nil {
		e.t.Fatal(err)
	}
	g := service.DefaultGeneratorSettings()
	g.DNS.PrimaryNS = "ns1.office.test"
	g.DNS.Hostmaster = "hostmaster@office.test"
	g.DNS.NameServers = []service.GeneratorNS{{Name: "ns1.office.test", Addresses: []string{"192.0.2.53"}}}
	g.DNS.DefaultTTL = ttl
	if _, err := e.svc.SetGeneratorSettings(ctx, service.SystemActor, g); err != nil {
		e.t.Fatal(err)
	}
}

func TestEnrollConnectAndBundle(t *testing.T) {
	e := newTestEnv(t)
	ctx := ctxT(t)

	// A server whose CA does not match the fingerprint is not trusted.
	other, _ := pki.Generate(nil)
	_, req := csr(t)
	if _, err := e.bootstrapClient(other.Fingerprint()).Enroll(ctx, &agentv1.EnrollRequest{Token: e.token("dns"), CsrDer: req,
		Role: agentv1.Role_ROLE_DNS}); err == nil {
		t.Fatal("enrollment with a wrong CA fingerprint succeeded")
	}

	a := e.enroll("dns")
	if id, err := agentpki.AgentID(a.cert); err != nil || id != a.id {
		t.Errorf("certificate agent ID %q %v", id, err)
	}
	if agentpki.Fingerprint(a.ca) != e.pki.Fingerprint() || !a.bundlePub.Equal(e.pki.BundlePublicKey()) {
		t.Error("enrollment returned foreign trust anchors")
	}
	if time.Until(a.cert.NotAfter) > pki.AgentCertLifetime || len(a.cert.ExtKeyUsage) != 1 || a.cert.ExtKeyUsage[0] != x509.ExtKeyUsageClientAuth {
		t.Errorf("client certificate %v %v", a.cert.NotAfter, a.cert.ExtKeyUsage)
	}

	st, err := a.client.GetDesiredState(ctx, &agentv1.GetDesiredStateRequest{})
	if err != nil || st.GetDesiredRevision() != 0 || st.GetMode() != agentv1.OperatingMode_OPERATING_MODE_AUTO_PULL ||
		st.GetPullInterval().AsDuration() != time.Minute {
		t.Fatalf("desired state %v %v", st, err)
	}

	stream, err := a.client.Connect(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if err := stream.Send(&agentv1.ConnectRequest{Message: &agentv1.ConnectRequest_Hello{Hello: &agentv1.Hello{
		AgentVersion: "1.0.1", Hostname: "ns1.office.test", ServiceVersion: "BIND 9.20.3",
		Health: &agentv1.HealthStatus{State: agentv1.HealthState_HEALTH_STATE_OK, CheckedAt: timestamppb.Now(),
			Checks: []*agentv1.HealthCheck{{Name: "rndc", State: agentv1.HealthState_HEALTH_STATE_OK}}},
	}}}); err != nil {
		t.Fatal(err)
	}
	msg, err := stream.Recv()
	if err != nil {
		t.Fatal(err)
	}
	w := msg.GetWelcome()
	if w == nil || w.GetDesiredRevision() != 0 || w.GetHeartbeatInterval().AsDuration() != 30*time.Second ||
		w.GetPullInterval().AsDuration() != time.Minute || w.GetMode() != agentv1.OperatingMode_OPERATING_MODE_AUTO_PULL {
		t.Fatalf("welcome %v", msg)
	}
	eventually(t, "connected state", func() bool { return e.hub.Connected(a.id) })
	ag, err := e.svc.GetAgent(ctx, a.id)
	if err != nil || !ag.Connected || ag.Hostname != "ns1.office.test" || ag.ServiceVersion != "BIND 9.20.3" ||
		ag.Health.State != "ok" || len(ag.Health.Checks) != 1 || ag.RemoteAddress != "127.0.0.1" {
		t.Errorf("agent after hello %+v %v", ag, err)
	}

	// A new revision is pushed to the connected agent.
	e.configureGenerator(3600)
	rev, created, err := e.svc.CreateRevision(ctx, service.SystemActor, "first")
	if err != nil || !created || rev.Number != 1 {
		t.Fatalf("create revision %+v %v %v", rev, created, err)
	}
	msg, err = stream.Recv()
	if err != nil {
		t.Fatal(err)
	}
	cmd := msg.GetCommand()
	if cmd.GetApplyRevision().GetRevision() != 1 || cmd.GetCommandId() == "" || cmd.GetApplyRevision().GetDeploymentId() == "" {
		t.Fatalf("command %v", msg)
	}
	if err := stream.Send(&agentv1.ConnectRequest{Message: &agentv1.ConnectRequest_CommandAck{CommandAck: &agentv1.CommandAck{
		CommandId: cmd.GetCommandId(), Accepted: true}}}); err != nil {
		t.Fatal(err)
	}

	// The bundle verifies with the key from the enrollment.
	b, err := a.client.FetchBundle(ctx, &agentv1.FetchBundleRequest{Revision: 1})
	if err != nil {
		t.Fatal(err)
	}
	contents := map[string][]byte{}
	for _, art := range b.GetArtifacts() {
		contents[art.GetPath()] = art.GetContent()
	}
	m, err := bundle.Verify(a.bundlePub, b.GetManifest(), b.GetSignature(), contents, bundle.Expect{AgentID: a.id, Role: "dns", MinRevision: 1})
	if err != nil {
		t.Fatalf("verify bundle: %v", err)
	}
	if m.Revision != 1 || len(m.Files) != len(rev.DNS) || !m.CreatedAt.Equal(rev.CreatedAt) {
		t.Errorf("manifest %+v", m)
	}
	if _, ok := contents["zones/db.office.test"]; !ok {
		t.Errorf("zone file missing: %v", m.Files)
	}
	again, err := a.client.FetchBundle(ctx, &agentv1.FetchBundleRequest{Revision: 1})
	if err != nil || !bytes.Equal(again.GetManifest(), b.GetManifest()) || !bytes.Equal(again.GetSignature(), b.GetSignature()) {
		t.Errorf("bundle is not deterministic: %v", err)
	}
	_, err = a.client.FetchBundle(ctx, &agentv1.FetchBundleRequest{Revision: 99})
	wantCode(t, err, codes.NotFound)

	// Heartbeat and deployment report.
	if err := stream.Send(&agentv1.ConnectRequest{Message: &agentv1.ConnectRequest_Heartbeat{Heartbeat: &agentv1.Heartbeat{
		AppliedRevision: 1, LastGoodRevision: 1,
		Health: &agentv1.HealthStatus{State: agentv1.HealthState_HEALTH_STATE_DEGRADED, Message: "slow"}}}}); err != nil {
		t.Fatal(err)
	}
	eventually(t, "heartbeat", func() bool {
		ag, err := e.svc.GetAgent(ctx, a.id)
		return err == nil && ag.AppliedRevision == 1 && ag.Health.State == "degraded"
	})
	start := time.Now().Add(-2 * time.Second)
	if _, err := a.client.ReportDeployment(ctx, &agentv1.ReportDeploymentRequest{
		DeploymentId: cmd.GetApplyRevision().GetDeploymentId(), Revision: 1,
		Trigger: agentv1.DeploymentTrigger_DEPLOYMENT_TRIGGER_PUSH, Outcome: agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_SUCCEEDED,
		ActiveRevision: 1, StartedAt: timestamppb.New(start), FinishedAt: timestamppb.Now(),
	}); err != nil {
		t.Fatal(err)
	}
	_, err = a.client.ReportDeployment(ctx, &agentv1.ReportDeploymentRequest{Revision: 1})
	wantCode(t, err, codes.InvalidArgument)
	ag, _ = e.svc.GetAgent(ctx, a.id)
	if ag.LastDeployment == nil || ag.LastDeployment.Trigger != "push" || ag.LastGoodRevision != 1 || ag.DesiredRevision != 1 {
		t.Errorf("agent after deployment %+v", ag)
	}

	// The agent closes its stream.
	if err := stream.CloseSend(); err != nil {
		t.Fatal(err)
	}
	if _, err := stream.Recv(); err != io.EOF { //nolint:errorlint // gRPC returns io.EOF as is
		t.Errorf("recv after close: %v", err)
	}
	eventually(t, "disconnect", func() bool { return !e.hub.Connected(a.id) })
}

func TestEnrollmentRejections(t *testing.T) {
	e := newTestEnv(t)
	ctx := ctxT(t)
	client := e.bootstrapClient(e.pki.Fingerprint())
	_, req := csr(t)

	tok := e.token("dns")
	_, err := client.Enroll(ctx, &agentv1.EnrollRequest{Token: tok, CsrDer: req, Role: agentv1.Role_ROLE_DHCP})
	wantCode(t, err, codes.PermissionDenied)
	_, err = client.Enroll(ctx, &agentv1.EnrollRequest{Token: tok, CsrDer: req})
	wantCode(t, err, codes.InvalidArgument)
	_, err = client.Enroll(ctx, &agentv1.EnrollRequest{Token: "garbage", CsrDer: req, Role: agentv1.Role_ROLE_DNS})
	wantCode(t, err, codes.Unauthenticated)
	_, err = client.Enroll(ctx, &agentv1.EnrollRequest{Token: tok, CsrDer: []byte("no csr"), Role: agentv1.Role_ROLE_DNS})
	wantCode(t, err, codes.InvalidArgument)

	// The failed attempts did not use up the token.
	if _, err := client.Enroll(ctx, &agentv1.EnrollRequest{Token: tok, CsrDer: req, Role: agentv1.Role_ROLE_DNS}); err != nil {
		t.Fatalf("enroll: %v", err)
	}
	_, err = client.Enroll(ctx, &agentv1.EnrollRequest{Token: tok, CsrDer: req, Role: agentv1.Role_ROLE_DNS})
	wantCode(t, err, codes.Unauthenticated)

	expired := e.token("dns")
	if _, err := e.st.Pool().Exec(ctx, `UPDATE enrollment_tokens SET created_at = now() - interval '2 hours',
		expires_at = now() - interval '1 hour' WHERE id = $1`, expired[:16]); err != nil {
		t.Fatal(err)
	}
	_, err = client.Enroll(ctx, &agentv1.EnrollRequest{Token: expired, CsrDer: req, Role: agentv1.Role_ROLE_DNS})
	wantCode(t, err, codes.Unauthenticated)
}

func TestEnrollmentRateLimit(t *testing.T) {
	e := newTestEnv(t)
	ctx := ctxT(t)
	client := e.bootstrapClient(e.pki.Fingerprint())
	for i := range 10 {
		_, err := client.Enroll(ctx, &agentv1.EnrollRequest{Token: "0123456789abcdef.wrong", Role: agentv1.Role_ROLE_DNS})
		if status.Code(err) != codes.Unauthenticated {
			t.Fatalf("attempt %d: %v", i, err)
		}
	}
	_, err := client.Enroll(ctx, &agentv1.EnrollRequest{Token: e.token("dns"), Role: agentv1.Role_ROLE_DNS})
	wantCode(t, err, codes.ResourceExhausted)
}

func TestClientCertificateRequired(t *testing.T) {
	e := newTestEnv(t)
	ctx := ctxT(t)
	a := e.enroll("dns")

	noCert := e.bootstrapClient(e.pki.Fingerprint())
	_, err := noCert.GetDesiredState(ctx, &agentv1.GetDesiredStateRequest{})
	wantCode(t, err, codes.Unauthenticated)
	_, err = noCert.FetchBundle(ctx, &agentv1.FetchBundleRequest{Revision: 1})
	wantCode(t, err, codes.Unauthenticated)
	stream, err := noCert.Connect(ctx)
	if err == nil {
		_, err = stream.Recv()
	}
	wantCode(t, err, codes.Unauthenticated)

	// Same agent ID, but issued by another CA.
	other, _ := pki.Generate(nil)
	key, req := csr(t)
	foreign, err := other.SignAgentCSR(req, a.id)
	if err != nil {
		t.Fatal(err)
	}
	forged := e.dial(agentpki.ClientTLSConfig(a.ca, tls.Certificate{Certificate: [][]byte{foreign.Raw}, PrivateKey: key}, serverName))
	_, err = forged.GetDesiredState(ctx, &agentv1.GetDesiredStateRequest{})
	wantCode(t, err, codes.Unauthenticated)

	// A certificate of our CA for an agent that does not exist.
	ghost, err := e.pki.SignAgentCSR(req, "00000000-0000-4000-8000-000000000001")
	if err != nil {
		t.Fatal(err)
	}
	ghostClient := e.dial(agentpki.ClientTLSConfig(a.ca, tls.Certificate{Certificate: [][]byte{ghost.Raw}, PrivateKey: key}, serverName))
	_, err = ghostClient.GetDesiredState(ctx, &agentv1.GetDesiredStateRequest{})
	wantCode(t, err, codes.Unauthenticated)
}

func TestRevocationClosesStream(t *testing.T) {
	e := newTestEnv(t)
	ctx := ctxT(t)
	a := e.enroll("dns")
	stream, err := a.client.Connect(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if err := stream.Send(&agentv1.ConnectRequest{Message: &agentv1.ConnectRequest_Hello{Hello: &agentv1.Hello{}}}); err != nil {
		t.Fatal(err)
	}
	if msg, err := stream.Recv(); err != nil || msg.GetWelcome() == nil {
		t.Fatalf("welcome %v %v", msg, err)
	}
	if _, err := e.svc.RevokeAgent(ctx, service.SystemActor, a.id); err != nil {
		t.Fatal(err)
	}
	_, err = stream.Recv()
	wantCode(t, err, codes.PermissionDenied)
	_, err = a.client.GetDesiredState(ctx, &agentv1.GetDesiredStateRequest{})
	wantCode(t, err, codes.PermissionDenied)
	_, req := csr(t)
	_, err = a.client.RenewCertificate(ctx, &agentv1.RenewCertificateRequest{CsrDer: req})
	wantCode(t, err, codes.PermissionDenied)
}

func TestFirstMessageMustBeHello(t *testing.T) {
	e := newTestEnv(t)
	ctx := ctxT(t)
	a := e.enroll("dns")
	stream, err := a.client.Connect(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if err := stream.Send(&agentv1.ConnectRequest{Message: &agentv1.ConnectRequest_Heartbeat{Heartbeat: &agentv1.Heartbeat{}}}); err != nil {
		t.Fatal(err)
	}
	_, err = stream.Recv()
	wantCode(t, err, codes.InvalidArgument)
}

func TestCertificateRenewal(t *testing.T) {
	e := newTestEnv(t)
	ctx := ctxT(t)
	day := 24 * time.Hour

	// Enrolled 15 days ago: the certificate expired yesterday.
	e.offset.Store(int64(-15 * day))
	old := e.enroll("dns")
	e.offset.Store(0)
	if !old.cert.NotAfter.Before(time.Now()) {
		t.Fatalf("certificate not expired: %v", old.cert.NotAfter)
	}
	_, err := old.client.GetDesiredState(ctx, &agentv1.GetDesiredStateRequest{})
	wantCode(t, err, codes.Unauthenticated)

	key, req := csr(t)
	resp, err := old.client.RenewCertificate(ctx, &agentv1.RenewCertificateRequest{CsrDer: req})
	if err != nil {
		t.Fatalf("renew expired certificate: %v", err)
	}
	renewed := e.agentFrom(old.id, "dns", key, resp.GetCertificateDer(), resp.GetCaCertificatesDer(), old.bundlePub)
	if !renewed.cert.NotAfter.After(time.Now().Add(13*day)) || renewed.cert.SerialNumber.Cmp(old.cert.SerialNumber) == 0 {
		t.Errorf("renewed certificate %v", renewed.cert.NotAfter)
	}
	if _, err := renewed.client.GetDesiredState(ctx, &agentv1.GetDesiredStateRequest{}); err != nil {
		t.Errorf("renewed certificate rejected: %v", err)
	}
	// The old certificate is worthless now, even for renewal.
	_, err = old.client.RenewCertificate(ctx, &agentv1.RenewCertificateRequest{CsrDer: req})
	wantCode(t, err, codes.Unauthenticated)

	// Renewal with a valid certificate replaces it as well.
	key2, req2 := csr(t)
	resp, err = renewed.client.RenewCertificate(ctx, &agentv1.RenewCertificateRequest{CsrDer: req2})
	if err != nil {
		t.Fatal(err)
	}
	third := e.agentFrom(old.id, "dns", key2, resp.GetCertificateDer(), resp.GetCaCertificatesDer(), old.bundlePub)
	_, err = renewed.client.GetDesiredState(ctx, &agentv1.GetDesiredStateRequest{})
	wantCode(t, err, codes.Unauthenticated)
	if _, err := third.client.GetDesiredState(ctx, &agentv1.GetDesiredStateRequest{}); err != nil {
		t.Errorf("third certificate rejected: %v", err)
	}
	_, err = third.client.RenewCertificate(ctx, &agentv1.RenewCertificateRequest{CsrDer: []byte("garbage")})
	wantCode(t, err, codes.InvalidArgument)

	// Expired for more than 30 days: a new enrollment is needed.
	e.offset.Store(int64(-45 * day))
	stale := e.enroll("dns")
	e.offset.Store(0)
	_, req3 := csr(t)
	_, err = stale.client.RenewCertificate(ctx, &agentv1.RenewCertificateRequest{CsrDer: req3})
	wantCode(t, err, codes.Unauthenticated)
}

func TestLeaseReports(t *testing.T) {
	e := newTestEnv(t)
	ctx := ctxT(t)
	dhcp := e.enroll("dhcp")
	dns := e.enroll("dns")

	lease := func(ip string, state agentv1.LeaseState) *agentv1.Lease {
		return &agentv1.Lease{Ip: ip, Mac: "00:11:22:33:44:55", Hostname: "pc", State: state,
			ExpiresAt: timestamppb.New(time.Now().Add(time.Hour)), SubnetId: 1}
	}
	report := func(seq uint64, full bool, ls ...*agentv1.Lease) uint64 {
		t.Helper()
		resp, err := dhcp.client.ReportLeases(ctx, &agentv1.ReportLeasesRequest{Sequence: seq, FullSync: full, Leases: ls})
		if err != nil {
			t.Fatalf("report %d: %v", seq, err)
		}
		return resp.GetAcknowledgedSequence()
	}
	count := func() int {
		t.Helper()
		page, err := e.svc.ListLeases(ctx, service.LeaseFilter{})
		if err != nil {
			t.Fatal(err)
		}
		return len(page.Items)
	}

	if ack := report(1, true, lease("192.0.2.10", agentv1.LeaseState_LEASE_STATE_ACTIVE), lease("192.0.2.11", agentv1.LeaseState_LEASE_STATE_ACTIVE)); ack != 1 || count() != 2 {
		t.Fatalf("full sync: ack %d, %d leases", ack, count())
	}
	// A repeated batch is dropped.
	if ack := report(1, true); ack != 1 || count() != 2 {
		t.Errorf("duplicate: ack %d, %d leases", ack, count())
	}
	if ack := report(3, false, lease("192.0.2.12", agentv1.LeaseState_LEASE_STATE_ACTIVE)); ack != 3 || count() != 3 {
		t.Errorf("delta: ack %d, %d leases", ack, count())
	}
	if ack := report(2, true); ack != 3 || count() != 3 {
		t.Errorf("old batch: ack %d, %d leases", ack, count())
	}
	if ack := report(4, true, lease("192.0.2.20", agentv1.LeaseState_LEASE_STATE_EXPIRED)); ack != 4 || count() != 1 {
		t.Errorf("second full sync: ack %d, %d leases", ack, count())
	}
	_, err := dhcp.client.ReportLeases(ctx, &agentv1.ReportLeasesRequest{Sequence: 0})
	wantCode(t, err, codes.InvalidArgument)
	_, err = dns.client.ReportLeases(ctx, &agentv1.ReportLeasesRequest{Sequence: 1, FullSync: true})
	wantCode(t, err, codes.PermissionDenied)
}

func TestShutdownClosesStreams(t *testing.T) {
	e := newTestEnv(t)
	ctx := ctxT(t)
	a := e.enroll("dns")
	stream, err := a.client.Connect(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if err := stream.Send(&agentv1.ConnectRequest{Message: &agentv1.ConnectRequest_Hello{Hello: &agentv1.Hello{}}}); err != nil {
		t.Fatal(err)
	}
	if _, err := stream.Recv(); err != nil {
		t.Fatal(err)
	}
	sctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	e.srv.Shutdown(sctx)
	_, err = stream.Recv()
	wantCode(t, err, codes.Unavailable)
}