feat: keep an accepted agent's public key after enrollment

internal/agent/handler_public_test.go

@@ -1590,10 +1590,14 @@ func (s *HandlerPublicTestSuite) TestStartInProgressKeepAlive() {
 			}
 
 			stop := agent.ExportStartInProgressKeepAlive(ctx, s.testAgent, mockMsg)
-			if !tt.cancelCtx {
-				// Let the ticker fire at least once before stopping.
-				time.Sleep(30 * time.Millisecond)
-			}
+
+			// Always let the keepalive goroutine reach its select before
+			// stop() closes the stop channel. Without this, a canceled
+			// context and a closed stop channel are both ready and select
+			// picks at random, so the ctx.Done() path is only sometimes
+			// taken -- which made its coverage flap between runs.
+			time.Sleep(30 * time.Millisecond)
+
 			stop()
 
 			tt.validateFunc(&calls)

internal/agent/pki/keypair.go

@@ -102,11 +102,23 @@ func (m *Manager) PrivateKey() ed25519.PrivateKey {
 // Fingerprint returns the SHA256 fingerprint of the public key in the
 // format "SHA256:<hex>". Returns an empty string if no public key is set.
 func (m *Manager) Fingerprint() string {
-	if len(m.publicKey) == 0 {
+	return FingerprintOf(m.publicKey)
+}
+
+// FingerprintOf returns the SHA256 fingerprint of the given public key in
+// the format "SHA256:<hex>". Returns an empty string when the key is empty,
+// so a record without a key never reports a fingerprint.
+//
+// Callers that hold a bare public key rather than a Manager use this, which
+// keeps one digest format across the codebase.
+func FingerprintOf(
+	pubKey ed25519.PublicKey,
+) string {
+	if len(pubKey) == 0 {
 		return ""
 	}
 
-	hash := sha256.Sum256(m.publicKey)
+	hash := sha256.Sum256(pubKey)
 
 	return "SHA256:" + hex.EncodeToString(hash[:])
 }

internal/controller/enrollment/accept.go

@@ -24,6 +24,7 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"strings"
 
 	"github.com/nats-io/nats.go/jetstream"
 
@@ -49,6 +50,14 @@ func (w *Watcher) AcceptAgent(
 		return fmt.Errorf("unmarshal pending agent %s: %w", machineID, err)
 	}
 
+	// Store the agent's key before anything is published: the store is the
+	// authority for every later verification, and acceptance is the only
+	// event allowed to write it. Deleting the pending entry below discards
+	// the only other copy.
+	if err := w.recordAgentKey(ctx, pending); err != nil {
+		return err
+	}
+
 	resp := pki.EnrollmentResponse{
 		Accepted:            true,
 		ControllerPublicKey: w.pkiProvider.PublicKey(),
@@ -117,6 +126,12 @@ func (w *Watcher) RejectAgent(
 		return fmt.Errorf("delete pending agent %s: %w", machineID, err)
 	}
 
+	// A rejected agent must not keep a stored key from an earlier
+	// acceptance, so nothing it signs verifies afterwards.
+	if err := w.RemoveAgentKey(ctx, machineID); err != nil {
+		return err
+	}
+
 	w.logger.Info(
 		"rejected agent enrollment",
 		slog.String("machine_id", machineID),
@@ -204,6 +219,13 @@ func (w *Watcher) findPendingBy(
 	}
 
 	for key := range lister.Keys() {
+		// Accepted agents' keys share this bucket and would otherwise
+		// unmarshal into a PendingAgent, making an already-accepted agent
+		// look pending.
+		if !strings.HasPrefix(key, kvPrefix) {
+			continue
+		}
+
 		entry, err := w.enrollmentKV.Get(ctx, key)
 		if err != nil {
 			continue

internal/controller/enrollment/export_test.go

@@ -78,6 +78,20 @@ func KVPrefix() string {
 	return kvPrefix
 }
 
+// AcceptedKVPrefix returns the acceptedKVPrefix constant for testing.
+func AcceptedKVPrefix() string {
+	return acceptedKVPrefix
+}
+
+// ExportRecordAgentKey exposes recordAgentKey for testing.
+func ExportRecordAgentKey(
+	ctx context.Context,
+	w *Watcher,
+	pending PendingAgent,
+) error {
+	return w.recordAgentKey(ctx, pending)
+}
+
 // EnrollSubject exposes enrollSubject for testing.
 func EnrollSubject(
 	namespace string,

internal/controller/enrollment/keystore.go

@@ -0,0 +1,192 @@
+// Copyright (c) 2026 John Dewey
+
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to
+// deal in the Software without restriction, including without limitation the
+// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+// sell copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+package enrollment
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/nats-io/nats.go/jetstream"
+
+	"github.com/osapi-io/osapi/internal/agent/pki"
+)
+
+// The three causes of a failed lookup stay distinct so an operator can tell
+// "this agent has not re-enrolled yet" from "something is forging messages"
+// from "the store could not be read". Collapsing them hides the difference
+// at exactly the moment it matters.
+var (
+	// ErrAgentKeyNotFound means no key is stored for the machine ID. During
+	// a rollout this is expected: no agent accepted before the store
+	// existed has one.
+	ErrAgentKeyNotFound = errors.New("no stored key for agent")
+
+	// ErrAgentKeyStoreUnavailable means the record could not be read or
+	// decoded. It never resembles "no stored key", and is never treated as
+	// verified.
+	ErrAgentKeyStoreUnavailable = errors.New("agent key store unavailable")
+
+	// ErrAgentKeySignatureMismatch means a key that is not this agent's
+	// signed the message.
+	ErrAgentKeySignatureMismatch = errors.New("agent signature mismatch")
+)
+
+// acceptedKey builds the KV key holding an accepted agent's record.
+func acceptedKey(
+	machineID string,
+) string {
+	return acceptedKVPrefix + machineID
+}
+
+// recordAgentKey stores an accepted agent's public key, keyed by machine ID,
+// so it outlives the pending record deleted on acceptance. The fingerprint is
+// recomputed from the key rather than copied from the request, because the
+// request's fingerprint is self-reported.
+//
+// Called only from the acceptance path.
+func (w *Watcher) recordAgentKey(
+	ctx context.Context,
+	pending PendingAgent,
+) error {
+	record := AcceptedAgent{
+		MachineID:   pending.MachineID,
+		Hostname:    pending.Hostname,
+		PublicKey:   pending.PublicKey,
+		Fingerprint: pki.FingerprintOf(pending.PublicKey),
+		AcceptedAt:  nowFn(),
+	}
+
+	data, err := marshalFn(record)
+	if err != nil {
+		return fmt.Errorf(
+			"marshal accepted agent %s: %w",
+			pending.MachineID,
+			err,
+		)
+	}
+
+	key := acceptedKey(pending.MachineID)
+	if _, err := w.enrollmentKV.Put(ctx, key, data); err != nil {
+		return fmt.Errorf(
+			"store accepted agent %s: %w",
+			pending.MachineID,
+			err,
+		)
+	}
+
+	w.invalidateAgentKey(pending.MachineID)
+
+	return nil
+}
+
+// LookupAgentKey returns the stored record for a machine ID. A missing
+// record and an unreadable store are different answers: the first is
+// ErrAgentKeyNotFound, the second ErrAgentKeyStoreUnavailable.
+func (w *Watcher) LookupAgentKey(
+	ctx context.Context,
+	machineID string,
+) (*AcceptedAgent, error) {
+	if record, ok := w.cachedAgentKey(machineID); ok {
+		return record, nil
+	}
+
+	entry, err := w.enrollmentKV.Get(ctx, acceptedKey(machineID))
+	if err != nil {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
+			return nil, fmt.Errorf("%w: %s", ErrAgentKeyNotFound, machineID)
+		}
+
+		return nil, errors.Join(
+			ErrAgentKeyStoreUnavailable,
+			fmt.Errorf("get accepted agent %s: %w", machineID, err),
+		)
+	}
+
+	var record AcceptedAgent
+	if err := unmarshalFn(entry.Value(), &record); err != nil {
+		return nil, errors.Join(
+			ErrAgentKeyStoreUnavailable,
+			fmt.Errorf("unmarshal accepted agent %s: %w", machineID, err),
+		)
+	}
+
+	w.cacheAgentKey(machineID, &record)
+
+	return &record, nil
+}
+
+// RemoveAgentKey deletes an agent's stored record. After it returns, nothing
+// signed by the removed key verifies.
+func (w *Watcher) RemoveAgentKey(
+	ctx context.Context,
+	machineID string,
+) error {
+	if err := w.enrollmentKV.Delete(ctx, acceptedKey(machineID)); err != nil {
+		return fmt.Errorf("remove accepted agent %s: %w", machineID, err)
+	}
+
+	w.invalidateAgentKey(machineID)
+
+	return nil
+}
+
+// cachedAgentKey returns a copy of the cached record for a machine ID, so a
+// caller cannot mutate what the next lookup returns.
+func (w *Watcher) cachedAgentKey(
+	machineID string,
+) (*AcceptedAgent, bool) {
+	w.keyCacheMu.RLock()
+	defer w.keyCacheMu.RUnlock()
+
+	record, ok := w.keyCache[machineID]
+	if !ok {
+		return nil, false
+	}
+
+	cached := *record
+
+	return &cached, true
+}
+
+// cacheAgentKey stores a copy of a record for later lookups. The cache is
+// invalidated by acceptance and removal, never by elapsed time: a time-based
+// entry would leave a removed agent verifying until it expired.
+func (w *Watcher) cacheAgentKey(
+	machineID string,
+	record *AcceptedAgent,
+) {
+	w.keyCacheMu.Lock()
+	defer w.keyCacheMu.Unlock()
+
+	cached := *record
+	w.keyCache[machineID] = &cached
+}
+
+// invalidateAgentKey drops a machine ID's cached record.
+func (w *Watcher) invalidateAgentKey(
+	machineID string,
+) {
+	w.keyCacheMu.Lock()
+	defer w.keyCacheMu.Unlock()
+
+	delete(w.keyCache, machineID)
+}