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) 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)