DDI/agent/core/store_test.go

// SPDX-License-Identifier: Apache-2.0

package core

import (
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"errors"
	"io/fs"
	"os"
	"path/filepath"
	"runtime"
	"slices"
	"strings"
	"testing"
	"time"

	"git.beugsiud.de/beug.k/DDI/internal/bundle"
)

func openTestStore(t *testing.T) *Store {
	t.Helper()
	s, err := OpenStore(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	return s
}

// storeTestRevision verifies and stores a signed test bundle.
func storeTestRevision(t *testing.T, s *Store, key ed25519.PrivateKey, rev uint64) Revision {
	t.Helper()
	b := signBundle(t, key, testAgentID, KindDNS, rev)
	contents := map[string][]byte{}
	for _, a := range b.Artifacts {
		contents[a.Path] = a.Content
	}
	m, err := bundle.Verify(key.Public().(ed25519.PublicKey), b.Manifest, b.Signature, contents,
		bundle.Expect{AgentID: testAgentID, Role: string(KindDNS)})
	if err != nil {
		t.Fatal(err)
	}
	r, err := s.StoreRevision(m, b.Manifest, b.Signature, contents)
	if err != nil {
		t.Fatal(err)
	}
	return r
}

func TestOpenStoreEmpty(t *testing.T) {
	if _, err := OpenStore(""); err == nil {
		t.Error("empty state directory accepted")
	}
}

func TestStoreRevision(t *testing.T) {
	s := openTestStore(t)
	_, key := newSigningKey(t)
	r := storeTestRevision(t, s, key, 7)
	if r.Number != 7 || r.Dir != filepath.Join(s.Dir(), "revisions", "7") || len(r.Files) != 2 {
		t.Fatalf("revision = %+v", r)
	}
	data, err := os.ReadFile(filepath.Join(r.Dir, "zones", "db.example.com"))
	if err != nil || !strings.Contains(string(data), "SOA ns 1 7") {
		t.Fatalf("artifact = %q, %v", data, err)
	}
	if runtime.GOOS != "windows" {
		fi, err := os.Stat(filepath.Join(r.Dir, "ddi-suite.conf"))
		if err != nil || fi.Mode().Perm() != 0o640 {
			t.Errorf("mode = %v, %v", fi.Mode(), err)
		}
	}
	loaded, err := s.LoadRevision(7)
	if err != nil || loaded.Number != 7 || loaded.Dir != r.Dir || len(loaded.Files) != 2 {
		t.Fatalf("LoadRevision = %+v, %v", loaded, err)
	}
	// Storing the same number again replaces the directory.
	if err := os.WriteFile(filepath.Join(r.Dir, "stale"), []byte("x"), 0o600); err != nil {
		t.Fatal(err)
	}
	storeTestRevision(t, s, key, 7)
	if _, err := os.Stat(filepath.Join(r.Dir, "stale")); !errors.Is(err, fs.ErrNotExist) {
		t.Errorf("old directory content survived: %v", err)
	}
	if _, err := s.LoadRevision(8); err == nil {
		t.Error("missing revision loaded")
	}
	entries, _ := os.ReadDir(filepath.Join(s.Dir(), "revisions"))
	for _, e := range entries {
		if strings.HasPrefix(e.Name(), ".") {
			t.Errorf("leftover %s", e.Name())
		}
	}
}

func TestStoreRevisionReservedPath(t *testing.T) {
	s := openTestStore(t)
	m := &bundle.Manifest{Revision: 1, Files: []bundle.File{{Path: "manifest.json", Mode: 0o644}}}
	if _, err := s.StoreRevision(m, []byte("{}"), []byte("sig"), map[string][]byte{"manifest.json": nil}); err == nil {
		t.Error("reserved artifact path accepted")
	}
	if revs, _ := s.ListRevisions(); len(revs) != 0 {
		t.Errorf("revisions = %v", revs)
	}
}

func TestLoadRevisionMismatch(t *testing.T) {
	s := openTestStore(t)
	_, key := newSigningKey(t)
	r := storeTestRevision(t, s, key, 3)
	if err := os.Rename(r.Dir, filepath.Join(s.Dir(), "revisions", "4")); err != nil {
		t.Fatal(err)
	}
	if _, err := s.LoadRevision(4); err == nil {
		t.Error("manifest of another revision accepted")
	}
}

func TestPrune(t *testing.T) {
	s := openTestStore(t)
	_, key := newSigningKey(t)
	for rev := uint64(1); rev <= 15; rev++ {
		storeTestRevision(t, s, key, rev)
	}
	for _, name := range []string{".tmp-123", ".old-5-abc", "notanumber", "007"} {
		if err := os.Mkdir(filepath.Join(s.Dir(), "revisions", name), 0o750); err != nil {
			t.Fatal(err)
		}
	}
	if err := s.Prune(KeepRevisions, 2, 3); err != nil {
		t.Fatal(err)
	}
	revs, err := s.ListRevisions()
	if err != nil {
		t.Fatal(err)
	}
	want := []uint64{2, 3, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}
	if !slices.Equal(revs, want) {
		t.Errorf("revisions = %v, want %v", revs, want)
	}
	for _, name := range []string{".tmp-123", ".old-5-abc"} {
		if _, err := os.Stat(filepath.Join(s.Dir(), "revisions", name)); !errors.Is(err, fs.ErrNotExist) {
			t.Errorf("%s not removed", name)
		}
	}
	if _, err := os.Stat(filepath.Join(s.Dir(), "revisions", "notanumber")); err != nil {
		t.Error("foreign directory removed")
	}
}

func TestStateRoundTrip(t *testing.T) {
	s := openTestStore(t)
	st, err := s.LoadState()
	if err != nil || st.Applied != 0 {
		t.Fatalf("empty state = %+v, %v", st, err)
	}
	now := time.Now().UTC().Truncate(time.Second)
	st = &State{
		Applied: 5, LastGood: 4, LastUnsuccessful: 5, LeaseSequence: 9,
		LastResult:     &DeploymentReport{Revision: 5, Outcome: OutcomeFailed, Trigger: TriggerPush, StartedAt: now, FinishedAt: now},
		PendingReports: []DeploymentReport{{Revision: 5, Outcome: OutcomeFailed, DeploymentID: "d"}},
	}
	if err := s.SaveState(st); err != nil {
		t.Fatal(err)
	}
	got, err := s.LoadState()
	if err != nil {
		t.Fatal(err)
	}
	if got.Applied != 5 || got.LastGood != 4 || got.LeaseSequence != 9 || got.LastResult.Outcome != OutcomeFailed ||
		len(got.PendingReports) != 1 || got.PendingReports[0].DeploymentID != "d" {
		t.Errorf("state = %+v", got)
	}
	// Human-readable JSON with the documented keys.
	raw, _ := os.ReadFile(filepath.Join(s.Dir(), "state.json"))
	var m map[string]any
	if err := json.Unmarshal(raw, &m); err != nil {
		t.Fatal(err)
	}
	for _, k := range []string{"applied_revision", "last_good_revision", "pending_reports", "lease_sequence"} {
		if _, ok := m[k]; !ok {
			t.Errorf("state.json lacks %s", k)
		}
	}
	entries, _ := os.ReadDir(s.Dir())
	for _, e := range entries {
		if strings.HasPrefix(e.Name(), ".tmp-") {
			t.Errorf("leftover temporary file %s", e.Name())
		}
	}
	if err := os.WriteFile(filepath.Join(s.Dir(), "state.json"), []byte("{broken"), 0o600); err != nil {
		t.Fatal(err)
	}
	if _, err := s.LoadState(); err == nil {
		t.Error("broken state.json accepted")
	}
}

func TestLock(t *testing.T) {
	s := openTestStore(t)
	unlock, err := s.Lock()
	if err != nil {
		t.Fatal(err)
	}
	if runtime.GOOS == "linux" {
		if _, err := s.Lock(); err == nil || !strings.Contains(err.Error(), "locked") {
			t.Errorf("second lock = %v", err)
		}
	}
	unlock()
	unlock2, err := s.Lock()
	if err != nil {
		t.Fatalf("lock after unlock: %v", err)
	}
	unlock2()
}

// testIdentity creates an identity signed by pki.
func testIdentity(t *testing.T, pki *testPKI) (*Identity, *ecdsa.PrivateKey) {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	cert, err := pki.issueClient(&key.PublicKey, testAgentID, time.Now().Add(-time.Hour), time.Now().Add(time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	pub, _ := newSigningKey(t)
	return &Identity{
		AgentID: testAgentID, Role: KindDNS, Server: "ddi.example.com:8444", ServerName: "ddi.example.com",
		EnrolledAt: time.Now().UTC().Truncate(time.Second), CA: pki.ca, Certificate: keyPair(key, cert), BundleKey: pub,
	}, key
}

func TestIdentityRoundTrip(t *testing.T) {
	s := openTestStore(t)
	if _, err := s.LoadIdentity(); !errors.Is(err, ErrNotEnrolled) {
		t.Fatalf("LoadIdentity = %v", err)
	}
	pki := newTestPKI(t)
	id, key := testIdentity(t, pki)
	if err := s.SaveIdentity(id, key, false); err != nil {
		t.Fatal(err)
	}
	if !s.Enrolled() {
		t.Error("not enrolled after SaveIdentity")
	}
	got, err := s.LoadIdentity()
	if err != nil {
		t.Fatal(err)
	}
	if got.AgentID != id.AgentID || got.Server != id.Server || got.ServerName != id.ServerName || got.Role != KindDNS ||
		!got.CA.Equal(pki.ca) || !got.BundleKey.Equal(id.BundleKey) || !got.Certificate.Leaf.Equal(id.Certificate.Leaf) ||
		!got.EnrolledAt.Equal(id.EnrolledAt) {
		t.Errorf("identity = %+v", got)
	}
	if runtime.GOOS != "windows" {
		fi, err := os.Stat(filepath.Join(s.Dir(), "identity", "key.pem"))
		if err != nil || fi.Mode().Perm() != 0o600 {
			t.Errorf("key.pem mode = %v, %v", fi.Mode(), err)
		}
		di, err := os.Stat(filepath.Join(s.Dir(), "identity"))
		if err != nil || di.Mode().Perm() != 0o700 {
			t.Errorf("identity dir mode = %v, %v", di.Mode(), err)
		}
	}
	// Already enrolled.
	if err := s.SaveIdentity(id, key, false); !errors.Is(err, errAlreadyEnrolled) {
		t.Errorf("second SaveIdentity = %v", err)
	}
	id2, key2 := testIdentity(t, pki)
	id2.Server = "other.example.com:8444"
	if err := s.SaveIdentity(id2, key2, true); err != nil {
		t.Fatal(err)
	}
	if got, err := s.LoadIdentity(); err != nil || got.Server != "other.example.com:8444" {
		t.Errorf("forced identity = %+v, %v", got, err)
	}
	entries, _ := os.ReadDir(s.Dir())
	for _, e := range entries {
		if strings.HasPrefix(e.Name(), ".") {
			t.Errorf("leftover %s", e.Name())
		}
	}
}

func TestSaveCertificateAndRecovery(t *testing.T) {
	s := openTestStore(t)
	pki := newTestPKI(t)
	id, key := testIdentity(t, pki)
	if err := s.SaveIdentity(id, key, false); err != nil {
		t.Fatal(err)
	}
	newKey, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	newCert, err := pki.issueClient(&newKey.PublicKey, testAgentID, time.Now().Add(-time.Minute), time.Now().Add(2*time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	if err := s.SaveCertificate(newKey, newCert); err != nil {
		t.Fatal(err)
	}
	got, err := s.LoadIdentity()
	if err != nil || !got.Certificate.Leaf.Equal(newCert) {
		t.Fatalf("renewed identity = %v", err)
	}

	// Simulate a crash between the two renames: key.pem is new, cert.pem
	// old, cert.pem.new holds the matching certificate.
	key3, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	cert3, err := pki.issueClient(&key3.PublicKey, testAgentID, time.Now().Add(-time.Minute), time.Now().Add(3*time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	keyPEM, _ := encodeKey(key3)
	dir := filepath.Join(s.Dir(), "identity")
	if err := os.WriteFile(filepath.Join(dir, "key.pem"), keyPEM, 0o600); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "cert.pem.new"), pemCert(cert3), 0o600); err != nil {
		t.Fatal(err)
	}
	got, err = s.LoadIdentity()
	if err != nil || !got.Certificate.Leaf.Equal(cert3) {
		t.Fatalf("recovered identity: %v", err)
	}
	if _, err := os.Stat(filepath.Join(dir, "cert.pem.new")); !errors.Is(err, fs.ErrNotExist) {
		t.Error("cert.pem.new left behind")
	}

	// A mismatch without a recoverable certificate is an error.
	key4, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	keyPEM, _ = encodeKey(key4)
	if err := os.WriteFile(filepath.Join(dir, "key.pem"), keyPEM, 0o600); err != nil {
		t.Fatal(err)
	}
	if _, err := s.LoadIdentity(); err == nil {
		t.Error("mismatched key pair accepted")
	}
}

func pemCert(c *x509.Certificate) []byte {
	return pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: c.Raw})
}

func TestCheckIssued(t *testing.T) {
	pki := newTestPKI(t)
	key, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	cert, err := pki.issueClient(&key.PublicKey, testAgentID, time.Now().Add(-time.Minute), time.Now().Add(time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	if _, err := checkIssued(cert.Raw, pki.ca, testAgentID, key); err != nil {
		t.Fatalf("checkIssued = %v", err)
	}
	if _, err := checkIssued(cert.Raw, pki.ca, "11111111-2222-4333-8444-555555555555", key); err == nil {
		t.Error("other agent ID accepted")
	}
	other, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if _, err := checkIssued(cert.Raw, pki.ca, testAgentID, other); err == nil {
		t.Error("certificate for another key accepted")
	}
	if _, err := checkIssued(cert.Raw, newTestPKI(t).ca, testAgentID, key); err == nil {
		t.Error("certificate of another CA accepted")
	}
	if _, err := checkIssued([]byte("junk"), pki.ca, testAgentID, key); err == nil {
		t.Error("junk accepted")
	}
	// A server clock slightly ahead is tolerated.
	future, err := pki.issueClient(&key.PublicKey, testAgentID, time.Now().Add(2*time.Minute), time.Now().Add(time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	if _, err := checkIssued(future.Raw, pki.ca, testAgentID, key); err != nil {
		t.Errorf("future NotBefore: %v", err)
	}
}