From 9020af832aa3b8e607c515c499128add21ba337d Mon Sep 17 00:00:00 2001 From: "Luis Gustavo S. Barreto" Date: Thu, 10 Sep 2026 12:14:52 -0300 Subject: [PATCH] fix(server): charge an install key use when a device is accepted used_times was incremented only where the server decided the admission itself during registration, so accepting from the review queue never spent a use and usage_limit had no effect on a manual-mode key. The charge moves into the device status transaction, which owns the transition to accepted. The store's increment is itself the limit check (it updates under usage_limit = 0 OR used_times < usage_limit), so an exhausted key surfaces as ErrInstallKeyExhausted and the accept fails rather than silently doing nothing. Running inside the transaction also rolls the use back when the status write fails, which is why the enrollment paths and the webhook callback no longer reserve and release it by hand, and why InstallKeyDecrementUsage goes away with them. used_times is not backfilled. An installation holding a manual key whose accepted devices already outnumber usage_limit keeps admitting until new accepts catch up with the limit. Backfilling from the accepted count would be truthful and would stop those keys at once, at the cost of turning an upgrade into a provisioning stop for whoever relied on the limit not working. Fixes: shellhub-io/shellhub#7068 --- server/api/services/device.go | 21 ++++++ server/api/services/device_test.go | 77 ++++++++++++++++++++++ server/api/services/enrollment.go | 21 ++---- server/api/services/enrollment_e2e_test.go | 51 +++++++++++++- server/api/services/errors.go | 1 + server/api/services/install-key.go | 17 +---- server/api/store/install_key.go | 6 -- server/api/store/mocks/mock_store.go | 57 ---------------- server/api/store/pg/install-key.go | 22 ------- 9 files changed, 155 insertions(+), 118 deletions(-) diff --git a/server/api/services/device.go b/server/api/services/device.go index 0a16b13d8ba..f6c5ddc2421 100644 --- a/server/api/services/device.go +++ b/server/api/services/device.go @@ -298,6 +298,23 @@ func (s *service) UpdateDeviceStatus(ctx context.Context, req *requests.DeviceUp return nil } +func (s *service) chargeInstallKeyUse(ctx context.Context, tenantID, installKeyID string) error { + if installKeyID == "" { + return nil + } + + key := &models.InstallKey{ID: installKeyID, TenantID: tenantID} + if err := s.store.InstallKeyIncrementUsage(ctx, key); err != nil { + if errors.Is(err, store.ErrNoDocuments) { + return ErrInstallKeyExhausted + } + + return err + } + + return nil +} + func (s *service) updateDeviceStatus(req *requests.DeviceUpdateStatus) store.TransactionCb { return func(ctx context.Context) error { namespace, err := s.store.NamespaceResolve(ctx, store.NamespaceTenantIDResolver, req.TenantID) @@ -382,6 +399,10 @@ func (s *service) updateDeviceStatus(req *requests.DeviceUpdateStatus) store.Tra return err } } + + if err := s.chargeInstallKeyUse(ctx, namespace.TenantID, device.InstallKeyID); err != nil { + return err + } } device.Status = newStatus diff --git a/server/api/services/device_test.go b/server/api/services/device_test.go index 47af1041cb1..9973a732d18 100644 --- a/server/api/services/device_test.go +++ b/server/api/services/device_test.go @@ -2371,6 +2371,83 @@ func TestUpdateDeviceStatus_licenseEvaluator(t *testing.T) { storeMock.AssertExpectations(t) } +func TestUpdateDeviceStatus_keylessDeviceSpendsNoKey(t *testing.T) { + envstest.SetEdition(t, envs.Community) + + savedHooks := deviceMergeHooks + deviceMergeHooks = nil + t.Cleanup(func() { deviceMergeHooks = savedHooks }) + + now := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC) + clockMock := clockmock.NewMockClock(t) + clockMock.On("Now").Return(now) + prevClockBackend := clock.DefaultBackend + t.Cleanup(func() { clock.DefaultBackend = prevClockBackend }) + clock.DefaultBackend = clockMock + + storeMock := storemock.NewMockStore(t) + queryOptionsMock := storemock.NewMockQueryOptions(t) + storeMock.On("Options").Return(queryOptionsMock).Maybe() + + ctx := context.Background() + const tenantID = "00000000-0000-0000-0000-000000000000" + + device := &models.Device{ + UID: "keyless", + Name: "keyless", + TenantID: tenantID, + Status: models.DeviceStatusPending, + Identity: &models.DeviceIdentity{MAC: "aa:bb:cc:dd:ee:ff"}, + } + accepted := *device + accepted.Status = models.DeviceStatusAccepted + accepted.StatusUpdatedAt = now + + storeMock. + On("NamespaceResolve", ctx, store.NamespaceTenantIDResolver, tenantID). + Return(&models.Namespace{TenantID: tenantID, MaxDevices: -1}, nil). + Once() + storeMock. + On("DeviceResolve", ctx, mock.Anything, store.DeviceUIDResolver, "keyless"). + Return(device, nil). + Once() + queryOptionsMock.On("WithDeviceStatus", models.DeviceStatusAccepted).Return(nil).Once() + storeMock. + On("DeviceResolve", ctx, mock.Anything, store.DeviceMACResolver, "aa:bb:cc:dd:ee:ff", mock.AnythingOfType("[]store.QueryOption")). + Return(nil, store.ErrNoDocuments). + Once() + storeMock. + On("DeviceResolve", ctx, mock.Anything, store.DeviceHostnameResolver, "keyless", mock.AnythingOfType("[]store.QueryOption")). + Return(nil, store.ErrNoDocuments). + Once() + storeMock.On("DeviceUpdate", ctx, &accepted).Return(nil).Once() + storeMock. + On("NamespaceIncrementDeviceCount", ctx, scope.MustBounded(tenantID), models.DeviceStatusPending, int64(-1)). + Return(nil). + Once() + storeMock. + On("NamespaceIncrementDeviceCount", ctx, scope.MustBounded(tenantID), models.DeviceStatusAccepted, int64(1)). + Return(nil). + Once() + storeMock. + On("InstallKeyEventStampDecision", ctx, scope.MustBounded(tenantID), "keyless", models.DeviceStatusAccepted, mock.Anything). + Return(nil). + Once() + storeMock. + On("WithTransaction", ctx, mock.AnythingOfType("store.TransactionCb")). + Return(func(ctx context.Context, cb store.TransactionCb) error { return cb(ctx) }). + Once() + + service := NewService(storeMock, privateKey, publicKey, storecache.NewNullCache()) + + require.NoError(t, service.UpdateDeviceStatus(ctx, &requests.DeviceUpdateStatus{ + TenantID: tenantID, UID: "keyless", Status: "accepted", + })) + + storeMock.AssertExpectations(t) + storeMock.AssertNotCalled(t, "InstallKeyIncrementUsage", mock.Anything, mock.Anything) +} + func TestDeviceUpdate(t *testing.T) { now := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC) storeMock := storemock.NewMockStore(t) diff --git a/server/api/services/enrollment.go b/server/api/services/enrollment.go index edb4ca1ab4f..22089deb042 100644 --- a/server/api/services/enrollment.go +++ b/server/api/services/enrollment.go @@ -123,29 +123,18 @@ func (s *service) applyEnrollmentDecision(ctx context.Context, decision enrollme switch decision { case enrollAccept: - if key != nil { - if err := s.store.InstallKeyIncrementUsage(ctx, key); err != nil { - log.WithError(err).WithField("install_key", key.Name).Warn("install key exhausted; device remains pending") - - return models.DeviceStatusPending - } - } - acceptReq := &requests.DeviceUpdateStatus{ TenantID: req.TenantID, UID: uid, Status: string(models.DeviceStatusAccepted), } if err := s.UpdateDeviceStatus(ctx, acceptReq); err != nil { - if key != nil { - if releaseErr := s.store.InstallKeyDecrementUsage(ctx, key); releaseErr != nil { - log.WithError(releaseErr).WithField("install_key", key.Name).Warn("failed to release reserved install key use") - } - } - - if errors.Is(err, ErrDeviceLicenseLimit) { + switch { + case errors.Is(err, ErrInstallKeyExhausted): + log.WithError(err).WithField("device_uid", uid).Warn("install key exhausted; device remains pending") + case errors.Is(err, ErrDeviceLicenseLimit): log.WithError(err).WithField("device_uid", uid).Warn("license limit reached; device remains pending") - } else { + default: log.WithError(err).WithField("device_uid", uid).Warn("auto-accept failed; device remains pending") } diff --git a/server/api/services/enrollment_e2e_test.go b/server/api/services/enrollment_e2e_test.go index 44353ae1612..69f83c2df9a 100644 --- a/server/api/services/enrollment_e2e_test.go +++ b/server/api/services/enrollment_e2e_test.go @@ -400,7 +400,7 @@ func TestEnrollmentE2E_CallbackSingleUse(t *testing.T) { } // TestEnrollmentE2E_CallbackHonorsKeyState proves the deferred-callback accept path mirrors the -// synchronous accept: it reserves a use against the key's limit and refuses once the key is no longer +// synchronous accept: it spends a use against the key's limit and refuses once the key is no longer // valid, so an outstanding token can't bypass the usage cap or accept with a key revoked after mint. func TestEnrollmentE2E_CallbackHonorsKeyState(t *testing.T) { e := setupEnrollmentE2E(t) @@ -432,7 +432,7 @@ func TestEnrollmentE2E_CallbackHonorsKeyState(t *testing.T) { return uid, callbackURL[strings.LastIndex(callbackURL, "/")+1:] } - t.Run("the callback accept reserves a use against the key limit", func(t *testing.T) { + t.Run("the callback accept spends a use against the key limit", func(t *testing.T) { uid, token := enrollDeferred(0x62, "webhook-limit", "aa:bb:cc:dd:ee:62") require.NoError(t, e.svc.ResolveEnrollmentCallback(context.Background(), &requests.EnrollmentCallback{Token: token, Decision: "accept"})) @@ -703,3 +703,50 @@ func TestEnrollmentE2E_HistoryCurrent(t *testing.T) { require.WithinDuration(t, firstDecidedAt, *older.DecidedAt, time.Second, "older event keeps the first accept time") require.WithinDuration(t, secondDecidedAt, *newer.DecidedAt, time.Second, "newer event keeps the second accept time") } + +// TestEnrollmentE2E_AcceptSpendsAUse covers the accepts a person makes: a use is spent when the +// device is admitted, not when it registers, and a key at its limit refuses the accept instead of +// admitting past the cap. +func TestEnrollmentE2E_AcceptSpendsAUse(t *testing.T) { + e := setupEnrollmentE2E(t) + + accept := func(uid string) error { + return e.svc.UpdateDeviceStatus(context.Background(), &requests.DeviceUpdateStatus{ + TenantID: e.tenantID, UID: uid, Status: string(models.DeviceStatusAccepted), + }) + } + + t.Run("manual review spends a use per accept and stops at the limit", func(t *testing.T) { + e.installKey(t, digest(0x90), "manual-capped", models.InstallKeyModeManual, models.InstallKeyTypeUser, func(k *models.InstallKey) { + k.UsageLimit = 1 + clearSecret(k) + }) + + first := e.enroll(t, "aa:bb:cc:dd:90:01", plaintextFor(0x90)) + second := e.enroll(t, "aa:bb:cc:dd:90:02", plaintextFor(0x90)) + require.Equal(t, models.DeviceStatusPending, e.status(t, first)) + require.Equal(t, models.DeviceStatusPending, e.status(t, second)) + require.Equal(t, 0, e.usedTimes(t, digest(0x90)), "a device waiting on a decision has spent nothing") + + require.NoError(t, accept(first)) + require.Equal(t, 1, e.usedTimes(t, digest(0x90))) + + require.ErrorIs(t, accept(second), ErrInstallKeyExhausted) + require.Equal(t, models.DeviceStatusPending, e.status(t, second), "a refused accept leaves the device in the queue") + require.Equal(t, 1, e.usedTimes(t, digest(0x90)), "a refused accept spends nothing") + }) + + t.Run("accepting a rejected device spends a use", func(t *testing.T) { + e.installKey(t, digest(0x91), "allow-capped", models.InstallKeyModeAllowlist, models.InstallKeyTypeUser, func(k *models.InstallKey) { + k.AllowedMACs = []string{"aa:bb:cc:dd:91:ff"} + clearSecret(k) + }) + + uid := e.enroll(t, "aa:bb:cc:dd:91:01", plaintextFor(0x91)) + require.Equal(t, models.DeviceStatusRejected, e.status(t, uid)) + require.Equal(t, 0, e.usedTimes(t, digest(0x91))) + + require.NoError(t, accept(uid)) + require.Equal(t, 1, e.usedTimes(t, digest(0x91))) + }) +} diff --git a/server/api/services/errors.go b/server/api/services/errors.go index b9271673b22..42e5d2b2a23 100644 --- a/server/api/services/errors.go +++ b/server/api/services/errors.go @@ -156,6 +156,7 @@ var ( ErrInstallKeyDuplicated = errors.New("InstallKey duplicated", ErrLayer, ErrCodeDuplicated) ErrInstallKeyForbidden = errors.New("the legacy install key cannot be modified", ErrLayer, ErrCodeForbidden) ErrInstallKeyInvalidField = errors.New("install key field is invalid", ErrLayer, ErrCodeInvalid) + ErrInstallKeyExhausted = errors.New("install key usage limit reached", ErrLayer, ErrCodeLimit) ErrAuthForbidden = errors.New("user is authenticated but cannot access this resource", ErrLayer, ErrCodeForbidden) ErrRoleForbidden = errors.New("role is forbidden", ErrLayer, ErrCodeForbidden) ErrUserDelete = errors.New("user couldn't be deleted", ErrLayer, ErrCodeInvalid) diff --git a/server/api/services/install-key.go b/server/api/services/install-key.go index 78b7754b94b..c7d08b8b629 100644 --- a/server/api/services/install-key.go +++ b/server/api/services/install-key.go @@ -21,7 +21,6 @@ import ( "github.com/shellhub-io/shellhub/pkg/models" "github.com/shellhub-io/shellhub/pkg/uuid" "github.com/shellhub-io/shellhub/server/api/store" - log "github.com/sirupsen/logrus" ) const ( @@ -509,21 +508,9 @@ func (s *service) ResolveEnrollmentCallback(ctx context.Context, req *requests.E return NewErrInstallKeyForbidden() } - if err := s.store.InstallKeyIncrementUsage(ctx, key); err != nil { - return NewErrInstallKeyForbidden() - } - - if err := s.UpdateDeviceStatus(ctx, &requests.DeviceUpdateStatus{ + return s.UpdateDeviceStatus(ctx, &requests.DeviceUpdateStatus{ TenantID: claims.TenantID, UID: claims.DeviceUID, Status: string(models.DeviceStatusAccepted), - }); err != nil { - if releaseErr := s.store.InstallKeyDecrementUsage(ctx, key); releaseErr != nil { - log.WithError(releaseErr).WithField("install_key", key.Name).Warn("failed to release reserved install key use") - } - - return err - } - - return nil + }) } diff --git a/server/api/store/install_key.go b/server/api/store/install_key.go index 93a4b1c0e4f..dabc727e5f1 100644 --- a/server/api/store/install_key.go +++ b/server/api/store/install_key.go @@ -59,12 +59,6 @@ type InstallKeyStore interface { // [ErrNoDocuments] when the key is already overused, closing the race between concurrent enrollments. InstallKeyIncrementUsage(ctx context.Context, installKey *models.InstallKey) (err error) - // InstallKeyDecrementUsage returns a use previously reserved by [InstallKeyIncrementUsage] when the - // enrollment it was reserved for did not go through (e.g. the accept failed), guarding at zero so a - // release never drives the counter negative. It returns [ErrNoDocuments] when there was nothing to - // release (counter already at zero). - InstallKeyDecrementUsage(ctx context.Context, installKey *models.InstallKey) (err error) - // InstallKeyEventCreate appends one immutable row to an install key's enrollment history. The store // stamps the event ID and timestamp. It returns an error, if any. InstallKeyEventCreate(ctx context.Context, event *models.InstallKeyEvent) (err error) diff --git a/server/api/store/mocks/mock_store.go b/server/api/store/mocks/mock_store.go index eea42a68c51..040eeaa6650 100644 --- a/server/api/store/mocks/mock_store.go +++ b/server/api/store/mocks/mock_store.go @@ -2435,63 +2435,6 @@ func (_c *MockStore_InstallKeyCreate_Call) RunAndReturn(run func(ctx context.Con return _c } -// InstallKeyDecrementUsage provides a mock function for the type MockStore -func (_mock *MockStore) InstallKeyDecrementUsage(ctx context.Context, installKey *models.InstallKey) error { - ret := _mock.Called(ctx, installKey) - - if len(ret) == 0 { - panic("no return value specified for InstallKeyDecrementUsage") - } - - var r0 error - if returnFunc, ok := ret.Get(0).(func(context.Context, *models.InstallKey) error); ok { - r0 = returnFunc(ctx, installKey) - } else { - r0 = ret.Error(0) - } - return r0 -} - -// MockStore_InstallKeyDecrementUsage_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'InstallKeyDecrementUsage' -type MockStore_InstallKeyDecrementUsage_Call struct { - *mock.Call -} - -// InstallKeyDecrementUsage is a helper method to define mock.On call -// - ctx context.Context -// - installKey *models.InstallKey -func (_e *MockStore_Expecter) InstallKeyDecrementUsage(ctx any, installKey any) *MockStore_InstallKeyDecrementUsage_Call { - return &MockStore_InstallKeyDecrementUsage_Call{Call: _e.mock.On("InstallKeyDecrementUsage", ctx, installKey)} -} - -func (_c *MockStore_InstallKeyDecrementUsage_Call) Run(run func(ctx context.Context, installKey *models.InstallKey)) *MockStore_InstallKeyDecrementUsage_Call { - _c.Call.Run(func(args mock.Arguments) { - var arg0 context.Context - if args[0] != nil { - arg0 = args[0].(context.Context) - } - var arg1 *models.InstallKey - if args[1] != nil { - arg1 = args[1].(*models.InstallKey) - } - run( - arg0, - arg1, - ) - }) - return _c -} - -func (_c *MockStore_InstallKeyDecrementUsage_Call) Return(err error) *MockStore_InstallKeyDecrementUsage_Call { - _c.Call.Return(err) - return _c -} - -func (_c *MockStore_InstallKeyDecrementUsage_Call) RunAndReturn(run func(ctx context.Context, installKey *models.InstallKey) error) *MockStore_InstallKeyDecrementUsage_Call { - _c.Call.Return(run) - return _c -} - // InstallKeyEventCreate provides a mock function for the type MockStore func (_mock *MockStore) InstallKeyEventCreate(ctx context.Context, event *models.InstallKeyEvent) error { ret := _mock.Called(ctx, event) diff --git a/server/api/store/pg/install-key.go b/server/api/store/pg/install-key.go index 787954db67e..8799cffa41b 100644 --- a/server/api/store/pg/install-key.go +++ b/server/api/store/pg/install-key.go @@ -191,28 +191,6 @@ func (pg *Pg) InstallKeyIncrementUsage(ctx context.Context, installKey *models.I return nil } -// InstallKeyDecrementUsage implements [store.InstallKeyStore]. -func (pg *Pg) InstallKeyDecrementUsage(ctx context.Context, installKey *models.InstallKey) error { - db := pg.GetConnection(ctx) - - r, err := db.NewUpdate(). - Model((*entity.InstallKey)(nil)). - Set("used_times = used_times - 1"). - Set("updated_at = ?", clock.Now()). - Where("key_digest = ? AND namespace_id = ?", installKey.ID, installKey.TenantID). - Where("used_times > 0"). - Exec(ctx) - if err != nil { - return fromSQLError(err) - } - - if rowsAffected, err := r.RowsAffected(); err != nil || rowsAffected == 0 { - return store.ErrNoDocuments - } - - return nil -} - // InstallKeyEventCreate implements [store.InstallKeyStore]. func (pg *Pg) InstallKeyEventCreate(ctx context.Context, event *models.InstallKeyEvent) error { db := pg.GetConnection(ctx)