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)