From 4bb4c5f3f9e3a55927dedce190c60ec38d3aeda7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D7=A0=CF=85=CE=B1=CE=B7=20=D7=A0=CF=85=CE=B1=CE=B7=D1=95?= =?UTF-8?q?=CF=83=CE=B7?= Date: Thu, 17 Sep 2026 22:13:46 -0700 Subject: [PATCH 1/2] feat: keep an accepted agent's public key after enrollment The controller had nowhere to keep an agent's public key: it lived only on the pending-enrollment record, which acceptance deletes. Every controller-side verification path therefore found no key and skipped, which is why the agent's response signing is currently decorative. Store the key in the enrollment bucket under an accepted. prefix, written only by acceptance and removed on rejection, with a per-machine-ID cache invalidated by both. A missing record, an unreadable store and a signature mismatch are three distinct sentinels, so an operator can tell "not enrolled yet" from "something is forging messages". Both scans of the bucket unmarshalled whatever keys they found as pending agents, which worked only because every key shared one prefix. An accepted record's JSON overlaps PendingAgent exactly, so it would have surfaced as a phantom pending agent; both scans now filter on the pending prefix. Verification itself lands next, on top of this store. Foundational phase of specs 002-agent-key-store. No behaviour changes while PKI is disabled. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01FuKUsHFG1EqZXamffh9M2c --- internal/agent/pki/keypair.go | 16 +- internal/controller/enrollment/accept.go | 22 ++ internal/controller/enrollment/export_test.go | 14 + internal/controller/enrollment/keystore.go | 192 ++++++++++ .../enrollment/keystore_public_test.go | 358 ++++++++++++++++++ .../enrollment/mocks/enrollment.gen.go | 45 ++- .../controller/enrollment/mocks/generate.go | 2 +- internal/controller/enrollment/types.go | 35 ++ internal/controller/enrollment/watcher.go | 20 + .../enrollment/watcher_public_test.go | 175 ++++++++- 10 files changed, 872 insertions(+), 7 deletions(-) create mode 100644 internal/controller/enrollment/keystore.go create mode 100644 internal/controller/enrollment/keystore_public_test.go diff --git a/internal/agent/pki/keypair.go b/internal/agent/pki/keypair.go index 983df48b1..b84dd5606 100644 --- a/internal/agent/pki/keypair.go +++ b/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:". 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:". 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[:]) } diff --git a/internal/controller/enrollment/accept.go b/internal/controller/enrollment/accept.go index 9d6353a50..095f1949d 100644 --- a/internal/controller/enrollment/accept.go +++ b/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 diff --git a/internal/controller/enrollment/export_test.go b/internal/controller/enrollment/export_test.go index eb37f472c..724f46e6d 100644 --- a/internal/controller/enrollment/export_test.go +++ b/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, diff --git a/internal/controller/enrollment/keystore.go b/internal/controller/enrollment/keystore.go new file mode 100644 index 000000000..54039ed61 --- /dev/null +++ b/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) +} diff --git a/internal/controller/enrollment/keystore_public_test.go b/internal/controller/enrollment/keystore_public_test.go new file mode 100644 index 000000000..81683514e --- /dev/null +++ b/internal/controller/enrollment/keystore_public_test.go @@ -0,0 +1,358 @@ +// 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_test + +import ( + "context" + "crypto/ed25519" + "encoding/json" + "errors" + "log/slog" + "testing" + "time" + + "github.com/nats-io/nats.go/jetstream" + "github.com/stretchr/testify/suite" + "go.uber.org/mock/gomock" + + "github.com/osapi-io/osapi/internal/agent/pki" + "github.com/osapi-io/osapi/internal/controller/enrollment" + enrollMocks "github.com/osapi-io/osapi/internal/controller/enrollment/mocks" + jobMocks "github.com/osapi-io/osapi/internal/job/mocks" +) + +type KeyStorePublicTestSuite struct { + suite.Suite + + ctx context.Context + mockCtrl *gomock.Controller + mockNC *enrollMocks.MockNATSSubscriber + mockKV *jobMocks.MockKeyValue + mockPKI *enrollMocks.MockPKIProvider + watcher *enrollment.Watcher + fixedTime time.Time + pubKey ed25519.PublicKey +} + +func (s *KeyStorePublicTestSuite) SetupTest() { + s.ctx = context.Background() + s.mockCtrl = gomock.NewController(s.T()) + s.mockNC = enrollMocks.NewMockNATSSubscriber(s.mockCtrl) + s.mockKV = jobMocks.NewMockKeyValue(s.mockCtrl) + s.mockPKI = enrollMocks.NewMockPKIProvider(s.mockCtrl) + s.fixedTime = time.Date(2026, 4, 11, 12, 0, 0, 0, time.UTC) + s.pubKey = make(ed25519.PublicKey, ed25519.PublicKeySize) + + enrollment.SetNowFn(func() time.Time { return s.fixedTime }) + + s.watcher = enrollment.NewWatcher( + slog.Default(), + s.mockNC, + s.mockKV, + s.mockPKI, + false, + "osapi", + ) +} + +func (s *KeyStorePublicTestSuite) TearDownTest() { + s.mockCtrl.Finish() +} + +func (s *KeyStorePublicTestSuite) SetupSubTest() { + // A fresh watcher per subtest. The key cache lives on the Watcher, so a + // record cached by an earlier case would satisfy a lookup that this case + // expects to reach the store. + s.watcher = enrollment.NewWatcher( + slog.Default(), + s.mockNC, + s.mockKV, + s.mockPKI, + false, + "osapi", + ) +} + +func (s *KeyStorePublicTestSuite) TearDownSubTest() { + enrollment.ResetNowFn() +} + +func (s *KeyStorePublicTestSuite) TestLookupAgentKey() { + tests := []struct { + name string + machineID string + setupMock func() + validateFunc func(*enrollment.AcceptedAgent, error) + }{ + { + name: "returns the stored record", + machineID: "machine-001", + setupMock: func() { + entry := jobMocks.NewMockKeyValueEntry(s.mockCtrl) + entry.EXPECT().Value().Return( + s.makeAcceptedJSON("machine-001", "web-01"), + ) + s.mockKV.EXPECT(). + Get(gomock.Any(), "accepted.machine-001"). + Return(entry, nil) + }, + validateFunc: func(record *enrollment.AcceptedAgent, err error) { + s.Require().NoError(err) + s.Equal("machine-001", record.MachineID) + s.Equal("web-01", record.Hostname) + s.Equal( + pki.FingerprintOf(s.pubKey), + record.Fingerprint, + ) + }, + }, + { + name: "reports a missing record as not found", + machineID: "machine-404", + setupMock: func() { + s.mockKV.EXPECT(). + Get(gomock.Any(), "accepted.machine-404"). + Return(nil, jetstream.ErrKeyNotFound) + }, + validateFunc: func(record *enrollment.AcceptedAgent, err error) { + s.Require().Error(err) + s.Nil(record) + s.Require().ErrorIs(err, enrollment.ErrAgentKeyNotFound) + s.NotErrorIs(err, enrollment.ErrAgentKeyStoreUnavailable) + }, + }, + { + name: "reports an unreadable store as unavailable", + machineID: "machine-001", + setupMock: func() { + s.mockKV.EXPECT(). + Get(gomock.Any(), "accepted.machine-001"). + Return(nil, errors.New("kv error")) + }, + validateFunc: func(record *enrollment.AcceptedAgent, err error) { + s.Require().Error(err) + s.Nil(record) + s.Require().ErrorIs(err, enrollment.ErrAgentKeyStoreUnavailable) + s.NotErrorIs(err, enrollment.ErrAgentKeyNotFound) + }, + }, + { + name: "reports an undecodable record as unavailable", + machineID: "machine-001", + setupMock: func() { + entry := jobMocks.NewMockKeyValueEntry(s.mockCtrl) + entry.EXPECT().Value().Return([]byte("bad json")) + s.mockKV.EXPECT(). + Get(gomock.Any(), "accepted.machine-001"). + Return(entry, nil) + }, + validateFunc: func(record *enrollment.AcceptedAgent, err error) { + s.Require().Error(err) + s.Nil(record) + s.Require().ErrorIs(err, enrollment.ErrAgentKeyStoreUnavailable) + s.NotErrorIs(err, enrollment.ErrAgentKeyNotFound) + }, + }, + } + + for _, tc := range tests { + s.Run(tc.name, func() { + tc.setupMock() + + tc.validateFunc(s.watcher.LookupAgentKey(s.ctx, tc.machineID)) + }) + } +} + +// TestLookupAgentKeyCaches proves the second lookup does not read the store: +// the Get is expected exactly once. +func (s *KeyStorePublicTestSuite) TestLookupAgentKeyCaches() { + entry := jobMocks.NewMockKeyValueEntry(s.mockCtrl) + entry.EXPECT().Value().Return(s.makeAcceptedJSON("machine-001", "web-01")) + s.mockKV.EXPECT(). + Get(gomock.Any(), "accepted.machine-001"). + Return(entry, nil). + Times(1) + + first, err := s.watcher.LookupAgentKey(s.ctx, "machine-001") + s.Require().NoError(err) + + second, err := s.watcher.LookupAgentKey(s.ctx, "machine-001") + s.Require().NoError(err) + s.Equal(first, second) +} + +// TestLookupAgentKeyReturnsCopy proves a caller cannot mutate what the next +// lookup returns, since the cached record is shared. +func (s *KeyStorePublicTestSuite) TestLookupAgentKeyReturnsCopy() { + entry := jobMocks.NewMockKeyValueEntry(s.mockCtrl) + entry.EXPECT().Value().Return(s.makeAcceptedJSON("machine-001", "web-01")) + s.mockKV.EXPECT(). + Get(gomock.Any(), "accepted.machine-001"). + Return(entry, nil). + Times(1) + + first, err := s.watcher.LookupAgentKey(s.ctx, "machine-001") + s.Require().NoError(err) + + first.Hostname = "attacker-01" + + second, err := s.watcher.LookupAgentKey(s.ctx, "machine-001") + s.Require().NoError(err) + s.Equal("web-01", second.Hostname) +} + +func (s *KeyStorePublicTestSuite) TestRemoveAgentKey() { + tests := []struct { + name string + machineID string + setupMock func() + validateFunc func(error) + }{ + { + name: "removes the stored record", + machineID: "machine-001", + setupMock: func() { + s.mockKV.EXPECT(). + Delete(gomock.Any(), "accepted.machine-001"). + Return(nil) + }, + validateFunc: func(err error) { + s.Require().NoError(err) + }, + }, + { + name: "returns error when the delete fails", + machineID: "machine-001", + setupMock: func() { + s.mockKV.EXPECT(). + Delete(gomock.Any(), "accepted.machine-001"). + Return(errors.New("delete error")) + }, + validateFunc: func(err error) { + s.Require().Error(err) + s.Contains(err.Error(), "remove accepted agent machine-001") + }, + }, + } + + for _, tc := range tests { + s.Run(tc.name, func() { + tc.setupMock() + + tc.validateFunc(s.watcher.RemoveAgentKey(s.ctx, tc.machineID)) + }) + } +} + +// TestRemoveAgentKeyInvalidatesCache proves a removed agent stops verifying +// immediately rather than when a cache entry would have expired: the store is +// read again after the removal. +func (s *KeyStorePublicTestSuite) TestRemoveAgentKeyInvalidatesCache() { + first := jobMocks.NewMockKeyValueEntry(s.mockCtrl) + first.EXPECT().Value().Return(s.makeAcceptedJSON("machine-001", "web-01")) + s.mockKV.EXPECT(). + Get(gomock.Any(), "accepted.machine-001"). + Return(first, nil) + + _, err := s.watcher.LookupAgentKey(s.ctx, "machine-001") + s.Require().NoError(err) + + s.mockKV.EXPECT(). + Delete(gomock.Any(), "accepted.machine-001"). + Return(nil) + s.Require().NoError(s.watcher.RemoveAgentKey(s.ctx, "machine-001")) + + s.mockKV.EXPECT(). + Get(gomock.Any(), "accepted.machine-001"). + Return(nil, jetstream.ErrKeyNotFound) + + _, err = s.watcher.LookupAgentKey(s.ctx, "machine-001") + s.Require().ErrorIs(err, enrollment.ErrAgentKeyNotFound) +} + +// TestRecordAgentKeyInvalidatesCache proves a re-acceptance is visible to the +// next lookup rather than shadowed by the previously cached record. +func (s *KeyStorePublicTestSuite) TestRecordAgentKeyInvalidatesCache() { + stale := jobMocks.NewMockKeyValueEntry(s.mockCtrl) + stale.EXPECT().Value().Return(s.makeAcceptedJSON("machine-001", "web-01")) + s.mockKV.EXPECT(). + Get(gomock.Any(), "accepted.machine-001"). + Return(stale, nil) + + _, err := s.watcher.LookupAgentKey(s.ctx, "machine-001") + s.Require().NoError(err) + + s.mockKV.EXPECT(). + Put(gomock.Any(), "accepted.machine-001", gomock.Any()). + Return(uint64(2), nil) + s.Require().NoError(enrollment.ExportRecordAgentKey( + s.ctx, + s.watcher, + enrollment.PendingAgent{ + MachineID: "machine-001", + Hostname: "web-01-renamed", + PublicKey: s.pubKey, + }, + )) + + fresh := jobMocks.NewMockKeyValueEntry(s.mockCtrl) + fresh.EXPECT().Value().Return( + s.makeAcceptedJSON("machine-001", "web-01-renamed"), + ) + s.mockKV.EXPECT(). + Get(gomock.Any(), "accepted.machine-001"). + Return(fresh, nil) + + record, err := s.watcher.LookupAgentKey(s.ctx, "machine-001") + s.Require().NoError(err) + s.Equal("web-01-renamed", record.Hostname) +} + +// makeAcceptedJSON creates serialized AcceptedAgent JSON. +func (s *KeyStorePublicTestSuite) makeAcceptedJSON( + machineID string, + hostname string, +) []byte { + s.T().Helper() + + record := enrollment.AcceptedAgent{ + MachineID: machineID, + Hostname: hostname, + PublicKey: s.pubKey, + Fingerprint: pki.FingerprintOf(s.pubKey), + AcceptedAt: s.fixedTime, + } + + data, err := json.Marshal(record) + s.Require().NoError(err) + + return data +} + +func TestKeyStorePublicTestSuite( + t *testing.T, +) { + // Not parallel: this package's tests swap package-level seams, and the + // rotation suite already runs in parallel. A second parallel suite races + // it over those globals. + suite.Run(t, new(KeyStorePublicTestSuite)) +} diff --git a/internal/controller/enrollment/mocks/enrollment.gen.go b/internal/controller/enrollment/mocks/enrollment.gen.go index 21dc9ed92..09432a860 100644 --- a/internal/controller/enrollment/mocks/enrollment.gen.go +++ b/internal/controller/enrollment/mocks/enrollment.gen.go @@ -1,19 +1,21 @@ // Code generated by MockGen. DO NOT EDIT. -// Source: github.com/osapi-io/osapi/internal/controller/enrollment (interfaces: NATSSubscriber,PKIProvider) +// Source: github.com/osapi-io/osapi/internal/controller/enrollment (interfaces: NATSSubscriber,PKIProvider,AgentKeyStore) // // Generated by this command: // -// mockgen -destination=enrollment.gen.go -package=mocks github.com/osapi-io/osapi/internal/controller/enrollment NATSSubscriber,PKIProvider +// mockgen -destination=enrollment.gen.go -package=mocks github.com/osapi-io/osapi/internal/controller/enrollment NATSSubscriber,PKIProvider,AgentKeyStore // // Package mocks is a generated GoMock package. package mocks import ( + context "context" ed25519 "crypto/ed25519" reflect "reflect" nats "github.com/nats-io/nats.go" + enrollment "github.com/osapi-io/osapi/internal/controller/enrollment" gomock "go.uber.org/mock/gomock" ) @@ -107,3 +109,42 @@ func (mr *MockPKIProviderMockRecorder) PublicKey() *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PublicKey", reflect.TypeOf((*MockPKIProvider)(nil).PublicKey)) } + +// MockAgentKeyStore is a mock of AgentKeyStore interface. +type MockAgentKeyStore struct { + ctrl *gomock.Controller + recorder *MockAgentKeyStoreMockRecorder + isgomock struct{} +} + +// MockAgentKeyStoreMockRecorder is the mock recorder for MockAgentKeyStore. +type MockAgentKeyStoreMockRecorder struct { + mock *MockAgentKeyStore +} + +// NewMockAgentKeyStore creates a new mock instance. +func NewMockAgentKeyStore(ctrl *gomock.Controller) *MockAgentKeyStore { + mock := &MockAgentKeyStore{ctrl: ctrl} + mock.recorder = &MockAgentKeyStoreMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockAgentKeyStore) EXPECT() *MockAgentKeyStoreMockRecorder { + return m.recorder +} + +// LookupAgentKey mocks base method. +func (m *MockAgentKeyStore) LookupAgentKey(ctx context.Context, machineID string) (*enrollment.AcceptedAgent, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "LookupAgentKey", ctx, machineID) + ret0, _ := ret[0].(*enrollment.AcceptedAgent) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// LookupAgentKey indicates an expected call of LookupAgentKey. +func (mr *MockAgentKeyStoreMockRecorder) LookupAgentKey(ctx, machineID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LookupAgentKey", reflect.TypeOf((*MockAgentKeyStore)(nil).LookupAgentKey), ctx, machineID) +} diff --git a/internal/controller/enrollment/mocks/generate.go b/internal/controller/enrollment/mocks/generate.go index 8db0fd238..6223f5c9c 100644 --- a/internal/controller/enrollment/mocks/generate.go +++ b/internal/controller/enrollment/mocks/generate.go @@ -21,4 +21,4 @@ // Package mocks provides generated mock implementations for testing. package mocks -//go:generate go tool go.uber.org/mock/mockgen -destination=enrollment.gen.go -package=mocks github.com/osapi-io/osapi/internal/controller/enrollment NATSSubscriber,PKIProvider +//go:generate go tool go.uber.org/mock/mockgen -destination=enrollment.gen.go -package=mocks github.com/osapi-io/osapi/internal/controller/enrollment NATSSubscriber,PKIProvider,AgentKeyStore diff --git a/internal/controller/enrollment/types.go b/internal/controller/enrollment/types.go index c256b0ed5..8ef15aa6d 100644 --- a/internal/controller/enrollment/types.go +++ b/internal/controller/enrollment/types.go @@ -23,6 +23,7 @@ package enrollment import ( + "context" "crypto/ed25519" "time" @@ -39,6 +40,40 @@ type PendingAgent struct { RequestedAt time.Time `json:"requested_at"` } +// AcceptedAgent is the stored record of an accepted agent's public key. It +// outlives the pending record, which is deleted on acceptance, so every +// later message from that agent can be verified against the key the agent +// enrolled with. +// +// Only enrollment acceptance creates or replaces a record: nothing an agent +// sends may introduce or change one. +type AcceptedAgent struct { + MachineID string `json:"machine_id"` + Hostname string `json:"hostname"` + PublicKey ed25519.PublicKey `json:"public_key"` + Fingerprint string `json:"fingerprint"` + AcceptedAt time.Time `json:"accepted_at"` + + // SupersededKey is the key replaced by the most recent rotation. It is + // absent unless a rotation is inside its grace period. + SupersededKey ed25519.PublicKey `json:"superseded_key,omitempty"` + + // SupersededUntil is the instant the superseded key stops being + // accepted. Storing the instant rather than a duration means a restart + // cannot silently extend the window. + SupersededUntil time.Time `json:"superseded_until,omitempty"` +} + +// AgentKeyStore is the narrow lookup the verification callers depend on, so +// the job client and target resolution need not import this package +// wholesale. Satisfied by *Watcher. +type AgentKeyStore interface { + LookupAgentKey( + ctx context.Context, + machineID string, + ) (*AcceptedAgent, error) +} + // NATSSubscriber defines the NATS operations needed by the enrollment // watcher for subscribing to enrollment requests and publishing responses. // Satisfied by the nats-client's *Client type (Subscribe + PublishCore). diff --git a/internal/controller/enrollment/watcher.go b/internal/controller/enrollment/watcher.go index 4d5880cd6..ecb5416ea 100644 --- a/internal/controller/enrollment/watcher.go +++ b/internal/controller/enrollment/watcher.go @@ -25,6 +25,8 @@ import ( "encoding/json" "fmt" "log/slog" + "strings" + "sync" "time" "github.com/nats-io/nats.go" @@ -45,6 +47,11 @@ var nowFn = time.Now // kvPrefix is the key prefix for pending enrollment entries. const kvPrefix = "enrollment." +// acceptedKVPrefix is the key prefix for accepted agents' stored keys. Both +// prefixes share the enrollment bucket, so every scan of pending entries +// must filter on kvPrefix rather than reading whatever keys it finds. +const acceptedKVPrefix = "accepted." + // Watcher monitors NATS for agent enrollment requests and manages // pending agents in a JetStream KV bucket. type Watcher struct { @@ -54,6 +61,12 @@ type Watcher struct { pkiProvider PKIProvider autoAccept bool namespace string + + // keyCache holds accepted agents' records by machine ID. Verification + // touches it for every response and heartbeat, so the lookup does not + // hit the KV in steady state. Invalidated by acceptance and removal. + keyCacheMu sync.RWMutex + keyCache map[string]*AcceptedAgent } // NewWatcher creates a new enrollment Watcher. @@ -72,6 +85,7 @@ func NewWatcher( pkiProvider: pkiProvider, autoAccept: autoAccept, namespace: namespace, + keyCache: make(map[string]*AcceptedAgent), } } @@ -191,6 +205,12 @@ func (w *Watcher) ListPending( var pending []PendingAgent for key := range lister.Keys() { + // The bucket also holds accepted agents' keys, which are not + // pending entries and would otherwise unmarshal into one. + if !strings.HasPrefix(key, kvPrefix) { + continue + } + entry, err := w.enrollmentKV.Get(ctx, key) if err != nil { w.logger.Warn( diff --git a/internal/controller/enrollment/watcher_public_test.go b/internal/controller/enrollment/watcher_public_test.go index 2147958c5..c13375619 100644 --- a/internal/controller/enrollment/watcher_public_test.go +++ b/internal/controller/enrollment/watcher_public_test.go @@ -209,6 +209,10 @@ func (s *WatcherPublicTestSuite) TestHandleEnrollmentRequestAutoAccept() { s.mockKV.EXPECT(). Delete(gomock.Any(), "enrollment.machine-001"). Return(nil) + + s.mockKV.EXPECT(). + Put(gomock.Any(), "accepted.machine-001", gomock.Any()). + Return(uint64(1), nil) }, msg: s.makeEnrollmentMsg("machine-001", "web-01", "SHA256:abc123"), }, @@ -264,6 +268,27 @@ func (s *WatcherPublicTestSuite) TestAcceptAgent() { s.mockKV.EXPECT(). Delete(gomock.Any(), "enrollment.machine-001"). Return(nil) + + s.mockKV.EXPECT(). + Put(gomock.Any(), "accepted.machine-001", gomock.Any()). + DoAndReturn(func( + _ context.Context, + _ string, + value []byte, + ) (uint64, error) { + var stored enrollment.AcceptedAgent + s.Require().NoError(json.Unmarshal(value, &stored)) + s.Equal("machine-001", stored.MachineID) + s.Equal("web-01", stored.Hostname) + s.Equal( + pki.FingerprintOf([]byte("test-pubkey")), + stored.Fingerprint, + "the fingerprint is recomputed, not taken from the request", + ) + s.Equal(s.fixedTime, stored.AcceptedAt) + + return uint64(1), nil + }) }, validateFunc: func(err error) { s.Require().NoError(err) @@ -309,9 +334,20 @@ func (s *WatcherPublicTestSuite) TestAcceptAgent() { Get(gomock.Any(), "enrollment.machine-001"). Return(mockEntry, nil) + s.mockKV.EXPECT(). + Put(gomock.Any(), "accepted.machine-001", gomock.Any()). + Return(uint64(1), nil) + s.mockPKI.EXPECT().PublicKey().Return(s.pubKey) - enrollment.SetMarshalFn(func(_ any) ([]byte, error) { + // The stored record marshals first; fail the response marshal. + calls := 0 + enrollment.SetMarshalFn(func(v any) ([]byte, error) { + calls++ + if calls == 1 { + return json.Marshal(v) + } + return nil, errors.New("marshal error") }) }, @@ -332,6 +368,10 @@ func (s *WatcherPublicTestSuite) TestAcceptAgent() { Get(gomock.Any(), "enrollment.machine-001"). Return(mockEntry, nil) + s.mockKV.EXPECT(). + Put(gomock.Any(), "accepted.machine-001", gomock.Any()). + Return(uint64(1), nil) + s.mockPKI.EXPECT().PublicKey().Return(s.pubKey) s.mockNC.EXPECT(). PublishCore("osapi.enroll.response.machine-001", gomock.Any()). @@ -354,6 +394,10 @@ func (s *WatcherPublicTestSuite) TestAcceptAgent() { Get(gomock.Any(), "enrollment.machine-001"). Return(mockEntry, nil) + s.mockKV.EXPECT(). + Put(gomock.Any(), "accepted.machine-001", gomock.Any()). + Return(uint64(1), nil) + s.mockPKI.EXPECT().PublicKey().Return(s.pubKey) s.mockNC.EXPECT(). PublishCore("osapi.enroll.response.machine-001", gomock.Any()). @@ -368,6 +412,48 @@ func (s *WatcherPublicTestSuite) TestAcceptAgent() { s.Contains(err.Error(), "delete pending agent machine-001") }, }, + { + name: "returns error when the key store rejects the record", + machineID: "machine-001", + setupMock: func() { + pendingData := s.makePendingJSON("machine-001", "web-01", "SHA256:abc123") + + mockEntry := jobMocks.NewMockKeyValueEntry(s.mockCtrl) + mockEntry.EXPECT().Value().Return(pendingData) + s.mockKV.EXPECT(). + Get(gomock.Any(), "enrollment.machine-001"). + Return(mockEntry, nil) + + s.mockKV.EXPECT(). + Put(gomock.Any(), "accepted.machine-001", gomock.Any()). + Return(uint64(0), errors.New("kv error")) + }, + validateFunc: func(err error) { + s.Require().Error(err) + s.Contains(err.Error(), "store accepted agent machine-001") + }, + }, + { + name: "returns error when the record fails to marshal", + machineID: "machine-001", + setupMock: func() { + pendingData := s.makePendingJSON("machine-001", "web-01", "SHA256:abc123") + + mockEntry := jobMocks.NewMockKeyValueEntry(s.mockCtrl) + mockEntry.EXPECT().Value().Return(pendingData) + s.mockKV.EXPECT(). + Get(gomock.Any(), "enrollment.machine-001"). + Return(mockEntry, nil) + + enrollment.SetMarshalFn(func(_ any) ([]byte, error) { + return nil, errors.New("marshal error") + }) + }, + validateFunc: func(err error) { + s.Require().Error(err) + s.Contains(err.Error(), "marshal accepted agent machine-001") + }, + }, } for _, tc := range tests { @@ -412,6 +498,10 @@ func (s *WatcherPublicTestSuite) TestRejectAgent() { s.mockKV.EXPECT(). Delete(gomock.Any(), "enrollment.machine-001"). Return(nil) + + s.mockKV.EXPECT(). + Delete(gomock.Any(), "accepted.machine-001"). + Return(nil) }, validateFunc: func(err error) { s.Require().NoError(err) @@ -517,6 +607,36 @@ func (s *WatcherPublicTestSuite) TestRejectAgent() { s.Contains(err.Error(), "delete pending agent machine-001") }, }, + { + name: "returns error when removing the stored key fails", + machineID: "machine-001", + reason: "denied", + setupMock: func() { + pendingData := s.makePendingJSON("machine-001", "web-01", "SHA256:abc123") + + mockEntry := jobMocks.NewMockKeyValueEntry(s.mockCtrl) + mockEntry.EXPECT().Value().Return(pendingData) + s.mockKV.EXPECT(). + Get(gomock.Any(), "enrollment.machine-001"). + Return(mockEntry, nil) + + s.mockNC.EXPECT(). + PublishCore("osapi.enroll.response.machine-001", gomock.Any()). + Return(nil) + + s.mockKV.EXPECT(). + Delete(gomock.Any(), "enrollment.machine-001"). + Return(nil) + + s.mockKV.EXPECT(). + Delete(gomock.Any(), "accepted.machine-001"). + Return(errors.New("delete error")) + }, + validateFunc: func(err error) { + s.Require().Error(err) + s.Contains(err.Error(), "remove accepted agent machine-001") + }, + }, } for _, tc := range tests { @@ -650,6 +770,34 @@ func (s *WatcherPublicTestSuite) TestListPending() { s.Len(pending, 0) }, }, + { + name: "skips accepted agent keys sharing the bucket", + setupMock: func() { + keys := make(chan string, 2) + keys <- "enrollment.machine-001" + keys <- "accepted.machine-001" + close(keys) + + mockLister := jobMocks.NewMockKeyLister(s.mockCtrl) + mockLister.EXPECT().Keys().Return(keys) + + s.mockKV.EXPECT(). + ListKeys(gomock.Any()). + Return(mockLister, nil) + + entry := jobMocks.NewMockKeyValueEntry(s.mockCtrl) + entry.EXPECT().Value().Return( + s.makePendingJSON("machine-001", "web-01", "SHA256:abc123"), + ) + s.mockKV.EXPECT(). + Get(gomock.Any(), "enrollment.machine-001"). + Return(entry, nil) + }, + validateFunc: func(pending []enrollment.PendingAgent, err error) { + s.Require().NoError(err) + s.Len(pending, 1) + }, + }, } for _, tc := range tests { @@ -706,6 +854,10 @@ func (s *WatcherPublicTestSuite) TestAcceptByHostname() { s.mockKV.EXPECT(). Delete(gomock.Any(), "enrollment.machine-001"). Return(nil) + + s.mockKV.EXPECT(). + Put(gomock.Any(), "accepted.machine-001", gomock.Any()). + Return(uint64(1), nil) }, validateFunc: func(err error) { s.Require().NoError(err) @@ -715,8 +867,9 @@ func (s *WatcherPublicTestSuite) TestAcceptByHostname() { name: "returns error when no matching hostname found", hostname: "nonexistent", setupMock: func() { - keys := make(chan string, 1) + keys := make(chan string, 2) keys <- "enrollment.machine-001" + keys <- "accepted.machine-001" close(keys) mockLister := jobMocks.NewMockKeyLister(s.mockCtrl) @@ -820,6 +973,10 @@ func (s *WatcherPublicTestSuite) TestAcceptByFingerprint() { s.mockKV.EXPECT(). Delete(gomock.Any(), "enrollment.machine-001"). Return(nil) + + s.mockKV.EXPECT(). + Put(gomock.Any(), "accepted.machine-001", gomock.Any()). + Return(uint64(1), nil) }, validateFunc: func(err error) { s.Require().NoError(err) @@ -922,6 +1079,10 @@ func (s *WatcherPublicTestSuite) TestAcceptByFingerprint() { s.mockKV.EXPECT(). Delete(gomock.Any(), "enrollment.machine-002"). Return(nil) + + s.mockKV.EXPECT(). + Put(gomock.Any(), "accepted.machine-002", gomock.Any()). + Return(uint64(1), nil) }, validateFunc: func(err error) { s.Require().NoError(err) @@ -973,6 +1134,10 @@ func (s *WatcherPublicTestSuite) TestAcceptByFingerprint() { s.mockKV.EXPECT(). Delete(gomock.Any(), "enrollment.machine-002"). Return(nil) + + s.mockKV.EXPECT(). + Put(gomock.Any(), "accepted.machine-002", gomock.Any()). + Return(uint64(1), nil) }, validateFunc: func(err error) { s.Require().NoError(err) @@ -1092,6 +1257,7 @@ func (s *WatcherPublicTestSuite) TestEnrollSubject() { func (s *WatcherPublicTestSuite) TestKVPrefix() { s.Equal("enrollment.", enrollment.KVPrefix()) + s.Equal("accepted.", enrollment.AcceptedKVPrefix()) } // makeEnrollmentMsg creates a *nats.Msg with a serialized EnrollmentRequest. @@ -1160,6 +1326,11 @@ func (s *WatcherPublicTestSuite) TestRejectByHostname() { s.mockKV.EXPECT(). Delete(gomock.Any(), "enrollment.machine-001"). Return(nil) + + // RemoveAgentKey: a rejected agent keeps no stored key. + s.mockKV.EXPECT(). + Delete(gomock.Any(), "accepted.machine-001"). + Return(nil) }, validateFunc: func(err error) { s.Require().NoError(err) From fcebe140d3f96b7299fc0f33ff2a5eb411f6fdf8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D7=A0=CF=85=CE=B1=CE=B7=20=D7=A0=CF=85=CE=B1=CE=B7=D1=95?= =?UTF-8?q?=CF=83=CE=B7?= Date: Thu, 17 Sep 2026 22:31:47 -0700 Subject: [PATCH 2/2] test: make the keepalive ctx.Done() path deterministic The canceled-context case closed the stop channel immediately after starting the keepalive goroutine, so the goroutine often reached its select with both stop and ctx.Done() already ready. Select picks uniformly at random among ready cases, so the ctx.Done() return ran only about half the time. That made coverage of that line flap between runs: it shows 0, 1 or 2 misses across recent commits on main with the file unchanged. Measured here at 4 of 10 runs missing it before this change, 0 of 10 after. The other three cases already slept before stopping; this drops the condition so every case does, which lets the goroutine reach its select while a canceled context is the only ready case. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01FuKUsHFG1EqZXamffh9M2c --- internal/agent/handler_public_test.go | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/internal/agent/handler_public_test.go b/internal/agent/handler_public_test.go index d8fd4ddf1..5b6ad5dbd 100644 --- a/internal/agent/handler_public_test.go +++ b/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)