fix(server): charge an install key use when a device is accepted

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

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)

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

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

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)

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

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)

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)