refactor(api): modernize device status update

- Refactor UpdateDeviceStatus to accept request object instead of
  individual parameters
- Add comprehensive logging with structured fields for better debugging
- Implement transaction-based device operations for data consistency
- Enhance device merging logic with proper error handling and validation
- Add extensive documentation for complex acceptance rules and
  environment-specific behaviors
- Improve billing integration with proper cloud environment detection
- Extract device limit validation and cloud billing handling into
  separate methods
- Update API routes to use modern request binding and validation
  patterns
- Add backward compatibility mapping for legacy status string values
- Enhance test coverage with realistic scenarios and proper mocking
- Remove deprecated adjustDeviceCounters utility in favor of inline
  operations
- Improve MAC address conflict detection and hostname validation
- Add proper transaction rollback support for failed operations

Author: heiytor

api/routes/device.go

@@ -215,28 +215,29 @@ func (h *Handler) LookupDevice(c gateway.Context) error {
 }
 
 func (h *Handler) UpdateDeviceStatus(c gateway.Context) error {
-	var req requests.DeviceUpdateStatus
-	if err := c.Bind(&req); err != nil {
-		return err
-	}
+	req := new(requests.DeviceUpdateStatus)
 
-	if err := c.Validate(&req); err != nil {
+	if err := c.Bind(req); err != nil {
 		return err
 	}
 
-	var tenant string
-	if c.Tenant() != nil {
-		tenant = c.Tenant().ID
+	// TODO: Remove this legacy status mapping in API v2.
+	// This mapping exists solely for backward compatibility with API consumers
+	// that were sending string values before the device status refactor.
+	status := map[string]string{
+		"accept":  string(models.DeviceStatusAccepted),
+		"reject":  string(models.DeviceStatusRejected),
+		"pending": string(models.DeviceStatusPending),
+		"unused":  string(models.DeviceStatusUnused),
 	}
 
-	status := map[string]models.DeviceStatus{
-		"accept":  models.DeviceStatusAccepted,
-		"reject":  models.DeviceStatusRejected,
-		"pending": models.DeviceStatusPending,
-		"unused":  models.DeviceStatusUnused,
+	req.Status = status[req.Status]
+
+	if err := c.Validate(req); err != nil {
+		return err
 	}
 
-	if err := h.service.UpdateDeviceStatus(c.Ctx(), tenant, models.UID(req.UID), status[req.Status]); err != nil {
+	if err := h.service.UpdateDeviceStatus(c.Ctx(), req); err != nil {
 		return err
 	}
 

api/services/device.go

@@ -11,6 +11,7 @@ import (
 	"github.com/shellhub-io/shellhub/pkg/clock"
 	"github.com/shellhub-io/shellhub/pkg/envs"
 	"github.com/shellhub-io/shellhub/pkg/models"
+	log "github.com/sirupsen/logrus"
 )
 
 const StatusAccepted = "accepted"
@@ -29,9 +30,27 @@ type DeviceService interface {
 	RenameDevice(ctx context.Context, uid models.UID, name, tenant string) error
 	LookupDevice(ctx context.Context, namespace, name string) (*models.Device, error)
 	OfflineDevice(ctx context.Context, uid models.UID) error
-	UpdateDeviceStatus(ctx context.Context, tenant string, uid models.UID, status models.DeviceStatus) error
 
 	UpdateDevice(ctx context.Context, req *requests.DeviceUpdate) error
+	// UpdateDeviceStatus updates a device's status. Devices that are already accepted cannot change their status.
+	//
+	// When accepting, if a device with the same MAC address is already accepted within the same namespace, it
+	// merges these devices unless a third device with the same hostname already exists and has a different MAC
+	// address. The merge transfers all sessions from the old device to the new one, renames the new device to
+	// preserve the old device's identity, and deletes the old device. Also, if another accepted device already
+	// uses the same hostname but has a different MAC address, the operation fails.
+	//
+	// Environment-specific Acceptance Rules:
+	//   - Community/Enterprise: Only checks the namespace's device limit
+	//   - Cloud (billing active): Reports device acceptance to billing service for quota/payment validation
+	//   - Cloud (billing inactive): Checks if the device is removed and evaluates namespace capabilities:
+	//     * If device was previously removed: removes from removed list, then evaluates billing
+	//     * If device was not removed: counts total removed devices and checks against limits, then evaluates billing
+	//     * Billing evaluation determines if the namespace can accept more devices based on subscription status
+	//
+	// All operations are performed within a database transaction to ensure consistency during device merging
+	// and counter updates.
+	UpdateDeviceStatus(ctx context.Context, req *requests.DeviceUpdateStatus) error
 }
 
 func (s *service) ListDevices(ctx context.Context, req *requests.DeviceList) ([]models.Device, int, error) {
@@ -217,138 +236,118 @@ func (s *service) OfflineDevice(ctx context.Context, uid models.UID) error {
 	return nil
 }
 
-// UpdateDeviceStatus updates the device status.
-func (s *service) UpdateDeviceStatus(ctx context.Context, tenant string, uid models.UID, status models.DeviceStatus) error {
-	namespace, err := s.store.NamespaceResolve(ctx, store.NamespaceTenantIDResolver, tenant)
-	if err != nil {
-		return NewErrNamespaceNotFound(tenant, err)
-	}
-
-	device, err := s.store.DeviceResolve(ctx, store.DeviceUIDResolver, string(uid), s.store.Options().InNamespace(tenant))
-	if err != nil {
-		return NewErrDeviceNotFound(uid, err)
-	}
-
-	if device.Status == models.DeviceStatusAccepted {
-		return NewErrDeviceStatusAccepted(nil)
-	}
-
-	// Store the original status for counter updates
-	originalStatus := device.Status
+func (s *service) UpdateDeviceStatus(ctx context.Context, req *requests.DeviceUpdateStatus) error {
+	return s.store.WithTransaction(ctx, s.updateDeviceStatus(req))
+}
 
-	// NOTICE: when the device is intended to be rejected or in pending status, we don't check for duplications as it
-	// is not going to be considered for connections.
-	if status == models.DeviceStatusPending || status == models.DeviceStatusRejected {
-		if err := s.store.DeviceUpdateStatus(ctx, uid, status); err != nil {
-			return err
+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)
+		if err != nil {
+			return NewErrNamespaceNotFound(req.TenantID, err)
 		}
 
-		if err := s.adjustDeviceCounters(ctx, tenant, originalStatus, status); err != nil { // nolint:revive
-			return err
+		device, err := s.store.DeviceResolve(ctx, store.DeviceUIDResolver, req.UID, s.store.Options().InNamespace(namespace.TenantID))
+		if err != nil {
+			return NewErrDeviceNotFound(models.UID(req.UID), err)
 		}
 
-		return nil
-	}
+		if device.Status == models.DeviceStatusAccepted {
+			log.WithFields(log.Fields{"device_uid": device.UID}).
+				Warn("cannot change status - device already accepted")
 
-	// NOTICE: when the intended status is not accepted, we return an error because these status are not allowed
-	// to be set by the user.
-	if status != models.DeviceStatusAccepted {
-		return NewErrDeviceStatusInvalid(string(status), nil)
-	}
+			return NewErrDeviceStatusAccepted(nil)
+		}
 
-	// NOTICE: when there is an already accepted device with the same MAC address, we need to update the device UID
-	// transfer the sessions and delete the old device.
-	sameMacDev, err := s.store.DeviceResolve(ctx, store.DeviceMACResolver, device.Identity.MAC, s.store.Options().WithDeviceStatus(models.DeviceStatusAccepted), s.store.Options().InNamespace(device.TenantID))
-	if err != nil && err != store.ErrNoDocuments {
-		return NewErrDeviceNotFound(models.UID(device.UID), err)
-	}
+		oldStatus := device.Status
+		newStatus := models.DeviceStatus(req.Status)
 
-	// TODO: move this logic to store's transactions.
-	if sameMacDev != nil && sameMacDev.UID != device.UID {
-		sameDevice, _ := s.store.DeviceResolve(ctx, store.DeviceHostnameResolver, device.Name, s.store.Options().WithDeviceStatus(models.DeviceStatusAccepted), s.store.Options().InNamespace(device.TenantID))
-		if sameDevice != nil && sameDevice.Identity.MAC != device.Identity.MAC {
-			return NewErrDeviceDuplicated(device.Name, nil)
+		if newStatus == device.Status {
+			return nil
 		}
 
-		if err := s.store.SessionUpdateDeviceUID(ctx, models.UID(sameMacDev.UID), models.UID(device.UID)); err != nil && err != store.ErrNoDocuments {
-			return err
-		}
+		if newStatus == models.DeviceStatusAccepted {
+			opts := []store.QueryOption{s.store.Options().WithDeviceStatus(models.DeviceStatusAccepted), s.store.Options().InNamespace(namespace.TenantID)}
+			existingMacDevice, err := s.store.DeviceResolve(ctx, store.DeviceMACResolver, device.Identity.MAC, opts...)
+			if err != nil && !errors.Is(err, store.ErrNoDocuments) {
+				log.WithError(err).
+					WithFields(log.Fields{"mac": device.Identity.MAC}).
+					Error("failed to retrieve device using MAC")
 
-		if err := s.store.DeviceRename(ctx, models.UID(device.UID), sameMacDev.Name); err != nil {
-			return err
-		}
+				return err
+			}
 
-		if err := s.store.DeviceDelete(ctx, models.UID(sameMacDev.UID)); err != nil {
-			return err
-		}
+			if existingMacDevice != nil && existingMacDevice.UID != device.UID {
+				existingNameDevice, err := s.store.DeviceResolve(ctx, store.DeviceHostnameResolver, device.Name, opts...)
+				if err != nil && !errors.Is(err, store.ErrNoDocuments) {
+					log.WithError(err).
+						WithFields(log.Fields{"name": device.Name}).
+						Error("failed to retrieve device using name")
 
-		// We need to decrease the accepted device count twice because we deleted the old device.
-		if err := s.store.NamespaceIncrementDeviceCount(ctx, tenant, models.DeviceStatusAccepted, -1); err != nil {
-			return err
-		}
+					return err
+				}
 
-		if err := s.store.DeviceUpdateStatus(ctx, uid, status); err != nil {
-			return err
-		}
+				if existingNameDevice != nil && existingNameDevice.Identity.MAC != device.Identity.MAC {
+					log.WithFields(log.Fields{"device_uid": device.UID, "device_mac": device.Identity.MAC, "conflicting_device_name": device.Name}).
+						Error("device merge blocked - hostname already used by device with different MAC address")
 
-		if err := s.adjustDeviceCounters(ctx, tenant, originalStatus, status); err != nil { // nolint:revive
-			return err
-		}
+					return NewErrDeviceDuplicated(device.Name, nil)
+				}
 
-		return nil
-	}
+				if err := s.mergeDevice(ctx, namespace.TenantID, existingMacDevice, device); err != nil {
+					log.WithError(err).
+						WithFields(log.Fields{"device_uid": device.UID, "existing_device_uid": existingMacDevice.UID, "device_mac": device.Identity.MAC}).
+						Error("device merge operation failed")
 
-	if sameDevice, _ := s.store.DeviceResolve(ctx, store.DeviceHostnameResolver, device.Name, s.store.Options().WithDeviceStatus(models.DeviceStatusAccepted), s.store.Options().InNamespace(device.TenantID)); sameDevice != nil {
-		return NewErrDeviceDuplicated(device.Name, nil)
-	}
+					return err
+				}
+			} else {
+				existingDevice, err := s.store.DeviceResolve(ctx, store.DeviceHostnameResolver, device.Name, opts...)
+				if err != nil && !errors.Is(err, store.ErrNoDocuments) {
+					log.WithError(err).
+						WithFields(log.Fields{"name": device.Name}).
+						Error("failed to retrieve device using name")
+
+					return err
+				}
 
-	if status != models.DeviceStatusAccepted {
-		if err := s.store.DeviceUpdateStatus(ctx, uid, status); err != nil {
-			return err
-		}
+				if existingDevice != nil {
+					log.WithFields(log.Fields{"device_uid": device.UID, "conflicting_device_name": device.Name}).
+						Error("device acceptance blocked - hostname already used by another device")
 
-		if err := s.adjustDeviceCounters(ctx, tenant, originalStatus, status); err != nil { // nolint:revive
-			return err
-		}
+					return NewErrDeviceDuplicated(device.Name, nil)
+				}
 
-		return nil
-	}
+				if err := s.checkDeviceLimits(ctx, namespace, device); err != nil {
+					log.WithError(err).WithFields(log.Fields{"device_uid": device.UID}).
+						Error("namespace's limit reached - cannot accept another device")
 
-	switch {
-	case envs.IsCloud() && envs.HasBilling():
-		if namespace.Billing.IsActive() {
-			if err := s.BillingReport(s.client, namespace.TenantID, ReportDeviceAccept); err != nil {
-				return NewErrBillingReportNamespaceDelete(err)
-			}
-		} else {
-			if device.Status != models.DeviceStatusRemoved && namespace.HasMaxDevices() && namespace.HasLimitDevicesReached() {
-				return NewErrDeviceRemovedFull(namespace.MaxDevices, nil)
-			}
+					return err
+				}
 
-			ok, err := s.BillingEvaluate(s.client, namespace.TenantID)
-			if err != nil {
-				return NewErrBillingEvaluate(err)
-			}
+				if envs.IsCloud() {
+					if err := s.handleCloudBilling(ctx, namespace); err != nil {
+						log.WithError(err).WithFields(log.Fields{"device_uid": device.UID, "billing_active": namespace.Billing.IsActive()}).
+							Error("billing validation failed")
 
-			if !ok {
-				return ErrDeviceLimit
+						return err
+					}
+				}
 			}
 		}
-	default:
-		if namespace.HasMaxDevices() && namespace.HasMaxDevicesReached() {
-			return NewErrDeviceMaxDevicesReached(namespace.MaxDevices)
+
+		if err := s.store.DeviceUpdate(ctx, namespace.TenantID, device.UID, &models.DeviceChanges{Status: newStatus}); err != nil {
+			return err
 		}
-	}
 
-	if err := s.store.DeviceUpdateStatus(ctx, uid, status); err != nil {
-		return err
-	}
+		for status, count := range map[models.DeviceStatus]int64{oldStatus: -1, newStatus: 1} {
+			if err := s.store.NamespaceIncrementDeviceCount(ctx, namespace.TenantID, status, count); err != nil {
+				return err
+			}
+		}
 
-	if err := s.adjustDeviceCounters(ctx, tenant, originalStatus, status); err != nil { // nolint:revive
-		return err
+		return nil
 	}
-
-	return nil
 }
 
 func (s *service) UpdateDevice(ctx context.Context, req *requests.DeviceUpdate) error {
@@ -371,3 +370,67 @@ func (s *service) UpdateDevice(ctx context.Context, req *requests.DeviceUpdate)
 
 	return s.store.DeviceUpdate(ctx, req.TenantID, req.UID, changes)
 }
+
+// mergeDevice merges an old device into a new device. It transfers all sessions from the old device to the new one and
+// renames the new device to preserve the old device's identity. The old device is then deleted and the namespace's device count is decremented.
+func (s *service) mergeDevice(ctx context.Context, tenantID string, oldDevice *models.Device, newDevice *models.Device) error {
+	// TODO: update tunnels as well?
+
+	if err := s.store.SessionUpdateDeviceUID(ctx, models.UID(oldDevice.UID), models.UID(newDevice.UID)); err != nil && !errors.Is(err, store.ErrNoDocuments) {
+		return err
+	}
+
+	if err := s.store.DeviceUpdate(ctx, tenantID, newDevice.UID, &models.DeviceChanges{Name: oldDevice.Name}); err != nil {
+		return err
+	}
+
+	if err := s.store.DeviceDelete(ctx, models.UID(oldDevice.UID)); err != nil {
+		return err
+	}
+
+	if err := s.store.NamespaceIncrementDeviceCount(ctx, tenantID, oldDevice.Status, -1); err != nil { //nolint:revive
+		return err
+	}
+
+	return nil
+}
+
+// checkDeviceLimits validates if the namespace can accept more devices based on environment-specific limits.
+func (s *service) checkDeviceLimits(ctx context.Context, namespace *models.Namespace, device *models.Device) error {
+	switch {
+	case envs.IsCloud():
+		if !namespace.Billing.IsActive() && device.Status != models.DeviceStatusRemoved &&
+			namespace.HasMaxDevices() && namespace.HasLimitDevicesReached() {
+
+			return NewErrDeviceRemovedFull(namespace.MaxDevices, nil)
+		}
+
+		return nil
+	default:
+		if namespace.HasMaxDevices() && namespace.HasMaxDevicesReached() {
+			return NewErrDeviceMaxDevicesReached(namespace.MaxDevices)
+		}
+
+		return nil
+	}
+}
+
+// handleCloudBilling processes billing-related operations for Cloud environment.
+// This function has side effects: it may delete removed devices and report to billing.
+func (s *service) handleCloudBilling(ctx context.Context, namespace *models.Namespace) error {
+	if namespace.Billing.IsActive() {
+		if err := s.BillingReport(s.client, namespace.TenantID, ReportDeviceAccept); err != nil {
+			return NewErrBillingReportNamespaceDelete(err)
+		}
+	} else {
+		ok, err := s.BillingEvaluate(s.client, namespace.TenantID)
+		switch {
+		case err != nil:
+			return NewErrBillingEvaluate(err)
+		case !ok:
+			return ErrDeviceLimit
+		}
+	}
+
+	return nil
+}